2615. Sum of Distances

Group equal values, then compute prefix and suffix sums of distances to avoid the O(n²) brute-force approach.

Problem statement

You are given a 0-indexed integer array nums. There exists an array arr of length nums.length, where arr[i] is the sum of |i - j| over all j such that nums[j] == nums[i] and j != i. If there is no such j, set arr[i] to be 0.

Return the array arr.

Examples

Example 1

Input: nums = [1,3,1,1,2]
Output: [5,0,3,4,0]

Explanation:

  • When i = 0, nums[0] == nums[2] and nums[0] == nums[3]. Therefore, arr[0] = |0 - 2| + |0 - 3| = 5.
  • When i = 1, arr[1] = 0 because there is no other index with value 3.
  • When i = 2, nums[2] == nums[0] and nums[2] == nums[3]. Therefore, arr[2] = |2 - 0| + |2 - 3| = 3.
  • When i = 3, nums[3] == nums[0] and nums[3] == nums[2]. Therefore, arr[3] = |3 - 0| + |3 - 2| = 4.
  • When i = 4, arr[4] = 0 because there is no other index with value 2.

Example 2

Input: nums = [0,5,3]
Output: [0,0,0]

Explanation: Since each element in nums is distinct, arr[i] = 0 for all i.

Constraints

  • 1 <= nums.length <= 10^5
  • 0 <= nums[i] <= 10^9

Explanation

The naive approach of checking every pair (i, j) would be O(n²), which is too slow for n <= 10^5. Instead, we process elements with the same value together and compute distances efficiently using a prefix-suffix technique.

For each value v in nums, let its occurrences be at indices [i₀, i₁, i₂, ..., iₖ]. For any occurrence at position iⱼ, we need:

  • Sum of distances to all occurrences to its left: Σ(iⱼ - iₓ) for x < j
  • Sum of distances to all occurrences to its right: Σ(iₓ - iⱼ) for x > j

Both sums can be computed incrementally. If we know the sum of distances for the previous occurrence and how many elements came before it, we can derive the next sum using the gap between consecutive indices.

Left-to-right pass (prefix)

For each index i, compute pre[i] = sum of |i - j| for all j < i where nums[j] == nums[i]. We maintain for each value:

  • mp — cumulative sum of distances so far
  • ind — last seen index
  • freq — count of occurrences

Right-to-left pass (suffix)

Similarly compute suff[i] = sum of |i - j| for all j > i with the same value.

Combine

ans[i] = pre[i] + suff[i] gives the total sum of distances to all identical elements.

Why it works

When processing an element at index i, if there are freq previous occurrences at index ind, the new cumulative distance is:

previous_sum + freq × |i - ind|

This accounts for the additional distance each previous element must stretch to reach this new position.

Complexity

  • Time: O(n) — two passes through the array with O(1) map operations
  • Space: O(n) — for prefix, suffix, and answer arrays, plus O(k) for maps where k is the number of distinct values

Code

#define ll long long

class Solution {
public:
    vector<long long> distance(vector<int>& nums) {
        int n = nums.size();
        map<int, ll> mp, ind, freq;
        vector<ll> pre(n), suff(n), ans(n);

        for (int i = 0; i < n; i++) {
            if (freq[nums[i]]) {
                pre[i] = mp[nums[i]] + freq[nums[i]] * abs(i - ind[nums[i]]);
            }
            ind[nums[i]] = i;
            mp[nums[i]] = pre[i];
            freq[nums[i]]++;
        }

        mp.clear(), ind.clear(), freq.clear();

        for (int i = n - 1; i >= 0; i--) {
            if (freq[nums[i]]) {
                suff[i] = mp[nums[i]] + freq[nums[i]] * abs(i - ind[nums[i]]);
            }
            freq[nums[i]]++;
            ind[nums[i]] = i;
            mp[nums[i]] = suff[i];
        }

        for (int i = 0; i < n; i++) {
            ans[i] = pre[i] + suff[i];
        }

        return ans;
    }
};