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]andnums[0] == nums[3]. Therefore,arr[0] = |0 - 2| + |0 - 3| = 5. - When
i = 1,arr[1] = 0because there is no other index with value3. - When
i = 2,nums[2] == nums[0]andnums[2] == nums[3]. Therefore,arr[2] = |2 - 0| + |2 - 3| = 3. - When
i = 3,nums[3] == nums[0]andnums[3] == nums[2]. Therefore,arr[3] = |3 - 0| + |3 - 2| = 4. - When
i = 4,arr[4] = 0because there is no other index with value2.
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^50 <= 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ₓ)forx < j - Sum of distances to all occurrences to its right:
Σ(iₓ - iⱼ)forx > 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 farind— last seen indexfreq— 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 withO(1)map operations - Space:
O(n)— for prefix, suffix, and answer arrays, plusO(k)for maps wherekis 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;
}
};