2615. Sum of Distances

Array
Hash Table
Prefix Sum

TIME

O(n)

SPACE

O(n)

Intuition

Group indices by value. For each group, use prefix sums to efficiently calculate the sum of distances for each index.

Solution

For sorted indices [i0, i1, ..., ik] of the same value, at position j:

Left contribution: j * indices[j] - prefix_sum[j]
Right contribution: (total_sum - prefix_sum[j+1]) - (count - j - 1) * indices[j]

# Group indices by value
groups = defaultdict(list)
for i, num in enumerate(nums):
    groups[num].append(i)

# For each group, calculate distances using prefix sum
for indices in groups.values():
    prefix = [0]
    for idx in indices:
        prefix.append(prefix[-1] + idx)
    
    n = len(indices)
    for j, idx in enumerate(indices):
        left = idx * j - prefix[j]
        right = (prefix[n] - prefix[j + 1]) - idx * (n - j - 1)
        arr[idx] = left + right

Why This Works

For index at position j in sorted group:

Left side: (idx - i0) + (idx - i1) + ... + (idx - i_{j-1}) = j * idx - (i0 + i1 + ... + i_{j-1}) = j * idx - prefix[j]
Right side: (i_{j+1} - idx) + ... + (ik - idx) = (sum of right) - (count of right) * idx

Complexity

Time: O(n)
Space: O(n)

Edge Cases

All distinct elements → all zeros
All same elements → calculate full distances
Single element → zero

Code

from typing import List
from collections import defaultdict


class Solution:
    def distance(self, nums: List[int]) -> List[int]:
        """
        2615. Sum of Distances
        Time: O(n)
        Space: O(n)
        """
        n = len(nums)
        arr = [0] * n

        # Group indices by value
        groups = defaultdict(list)
        for i, num in enumerate(nums):
            groups[num].append(i)

        # For each group, calculate distances using prefix sum
        for indices in groups.values():
            # Build prefix sum of indices
            prefix = [0]
            for idx in indices:
                prefix.append(prefix[-1] + idx)

            count = len(indices)
            for j, idx in enumerate(indices):
                # Left contribution: j * idx - sum of indices to the left
                left = idx * j - prefix[j]
                # Right contribution: sum of indices to the right - (count - j - 1) * idx
                right = (prefix[count] - prefix[j + 1]) - idx * (count - j - 1)
                arr[idx] = left + right

        return arr

import java.util.*;

/**
 * 2615. Sum of Distances
 * Time: O(n)
 * Space: O(n)
 */
class Solution {
    public long[] distance(int[] nums) {
        int n = nums.length;
        long[] arr = new long[n];

        // Group indices by value
        Map<Integer, List<Integer>> groups = new HashMap<>();
        for (int i = 0; i < n; i++) {
            groups.computeIfAbsent(nums[i], k -> new ArrayList<>()).add(i);
        }

        // For each group, calculate distances using prefix sum
        for (List<Integer> indices : groups.values()) {
            int count = indices.size();

            // Build prefix sum of indices
            long[] prefix = new long[count + 1];
            for (int j = 0; j < count; j++) {
                prefix[j + 1] = prefix[j] + indices.get(j);
            }

            for (int j = 0; j < count; j++) {
                long idx = indices.get(j);
                // Left contribution
                long left = idx * j - prefix[j];
                // Right contribution
                long right = (prefix[count] - prefix[j + 1]) - idx * (count - j - 1);
                arr[(int) idx] = left + right;
            }
        }

        return arr;
    }
}