Sum of Distances
MedPrompt
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.
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 <= 1050 <= nums[i] <= 109
Note: This question is the same as 2121: Intervals Between Identical Elements.
Approaches
3 approaches with complexity analysis and trade-offs.
The brute-force approach is the most straightforward way to solve the problem. It directly translates the problem description into code. We iterate through each element of the array and, for each element, we perform another full scan of the array to find all other elements with the same value. When a match is found at a different index, we calculate the absolute difference of the indices and add it to a running sum for the current element.
Algorithm
- Initialize a result array
arrof sizen(length ofnums) with zeros. - Iterate through the input array
numswith an indexifrom0ton-1. - For each
i, initialize a variablecurrentSumto0. - Start a nested loop with an index
jfrom0ton-1. - Inside the nested loop, check if
iis not equal tojand ifnums[i]is equal tonums[j]. - If both conditions are true, calculate the absolute difference
|i - j|and add it tocurrentSum. - After the inner loop completes, assign the value of
currentSumtoarr[i]. - After the outer loop completes, return the
arr.
Walkthrough
This method involves using two nested loops. The outer loop iterates through each index i of the nums array, for which we want to calculate the sum of distances. The inner loop iterates through every other index j in the array. Inside the inner loop, we check if nums[j] is equal to nums[i]. If they are equal and i is not the same as j, we compute |i - j| and add it to a temporary sum. This sum is accumulated for all such j's. Once the inner loop finishes, the accumulated sum is the final value for arr[i]. This process is repeated for all indices i.
class Solution { public long[] distance(int[] nums) { int n = nums.length; long[] arr = new long[n]; for (int i = 0; i < n; i++) { long currentSum = 0; for (int j = 0; j < n; j++) { if (i == j) { continue; } if (nums[i] == nums[j]) { currentSum += Math.abs(i - j); } } arr[i] = currentSum; } return arr; }}Complexity
Time
O(N^2) - For each of the N elements in the input array, we iterate through the entire array again. This results in N * N operations.
Space
O(N) - We need an array of size N to store the results. If the output array is not considered, the space complexity is O(1).
Trade-offs
Pros
Simple to understand and implement.
Requires minimal auxiliary data structures.
Cons
Extremely inefficient for large inputs due to its quadratic time complexity.
Will result in a 'Time Limit Exceeded' (TLE) error on most competitive programming platforms for the given constraints.
A better approach than brute-force is to first group all the indices based on the value at those indices. By doing this, we avoid scanning the entire array for each element. We can use a hash map to store the lists of indices for each unique number. Then, for each group of indices, we can perform the distance calculation. This confines the distance calculation to only the relevant indices.
Algorithm
- Create a
HashMapwhere keys are the unique numbers fromnumsand values areList<Integer>of indices where those numbers appear. - Iterate through
numsfromi = 0ton-1and populate the map. For eachnums[i], add the indexito the corresponding list in the map. - Initialize a result array
arrof sizen. - Iterate through each entry in the map.
- For each list of indices in the map:
a. If the list has one or zero elements, the corresponding distances are 0, so we can skip.
b. For each index
iin the list, calculate its sum of distances by iterating through all other indicesjin the same list and summing up|i - j|. c. Store this sum inarr[i]. - Return
arr.
Walkthrough
We first preprocess the array to build a map from each number to a list of indices where it appears. For example, if nums = [1,3,1,1,2], the map would be {1: [0, 2, 3], 3: [1], 2: [4]}. After building this map, we iterate through its values (the lists of indices). For each list, and for each index i within that list, we calculate the sum of |i - j| for all other indices j in the same list. This is still a nested loop, but it's performed on smaller lists instead of the entire nums array.
import java.util.*; class Solution { public long[] distance(int[] nums) { int n = nums.length; long[] arr = new long[n]; Map<Integer, List<Integer>> map = new HashMap<>(); for (int i = 0; i < n; i++) { map.computeIfAbsent(nums[i], k -> new ArrayList<>()).add(i); } for (List<Integer> indices : map.values()) { if (indices.size() <= 1) { continue; } for (int i = 0; i < indices.size(); i++) { long currentSum = 0; int currentIndex = indices.get(i); for (int j = 0; j < indices.size(); j++) { if (i == j) { continue; } currentSum += Math.abs(currentIndex - indices.get(j)); } arr[currentIndex] = currentSum; } } return arr; }}Complexity
Time
O(N + Σ(k_v^2)) - where N is for building the map and k_v is the number of occurrences of value v. In the worst case, if all elements are the same, k_v = N, leading to O(N^2) complexity. In the best case (all unique elements), it's O(N).
Space
O(N) - In the worst case, all elements are unique, and the map will store N keys, each with a list of one index. We also need O(N) space for the result array.
Trade-offs
Pros
More efficient than pure brute-force, especially if the number of occurrences for any single value is small.
Organizes the problem by grouping relevant data together.
Cons
The time complexity is still quadratic in the worst-case scenario (when all elements in the input array are the same).
Can be inefficient if there are many occurrences of the same number.
The most efficient approach involves calculating the sum of distances in two separate passes using the concept of prefix and suffix sums. The total sum for an index i is the sum of distances to identical elements on its left plus the sum of distances to identical elements on its right. We can calculate all left-side sums in one pass (left-to-right) and all right-side sums in a second pass (right-to-left).
Algorithm
- The total sum of distances for an index
ican be split into a left part and a right part:arr[i] = (sum of distances to identical elements on the left) + (sum of distances to identical elements on the right)arr[i] = (count_left * i - sum_indices_left) + (sum_indices_right - count_right * i) - First Pass (Left-to-Right):
a. Initialize a result array
arr(long), acountMap, and asumMap. b. Iterateifrom0ton-1. Fornums[i], get thecountandsumof indices of its previous occurrences from the maps. c. Calculate the left-side sum:arr[i] = (long)count * i - sum. d. Update the maps with the current indexi: increment count and addito the sum fornums[i]. - Second Pass (Right-to-Left):
a. Clear or re-initialize the
countMapandsumMap. b. Iterateifromn-1down to0. Fornums[i], get thecountandsumof indices of its occurrences to the right (which have been processed in this reverse pass). c. Calculate the right-side sum:(long)sum - (long)count * i, and add it to the existingarr[i]. d. Update the maps with the current indexi. - Return
arr.
Walkthrough
This method cleverly breaks down the calculation. For any index i, the sum of distances to other identical elements j is Σ|i-j|. This can be split into Σ(i-j) for j < i and Σ(j-i) for j > i.
- The left sum
Σ(i-j)can be rewritten ascount_left * i - Σj(wherecount_leftis the number of identical elements to the left andΣjis the sum of their indices). - The right sum
Σ(j-i)can be rewritten asΣj - count_right * i.
We can compute these values efficiently. In a first pass from left to right, we use maps to track the running count and sum of indices for each number. This allows us to compute the left-side sum for each element. In a second pass from right to left, we do the same to compute the right-side sum and add it to our result. This avoids any nested loops and achieves a linear time solution.
import java.util.*; class Solution { public long[] distance(int[] nums) { int n = nums.length; long[] arr = new long[n]; // Pass 1: Left to Right Map<Integer, Integer> countMap = new HashMap<>(); Map<Integer, Long> sumMap = new HashMap<>(); for (int i = 0; i < n; i++) { int num = nums[i]; int count = countMap.getOrDefault(num, 0); long sum = sumMap.getOrDefault(num, 0L); arr[i] = (long)count * i - sum; countMap.put(num, count + 1); sumMap.put(num, sum + i); } // Pass 2: Right to Left countMap.clear(); sumMap.clear(); for (int i = n - 1; i >= 0; i--) { int num = nums[i]; int count = countMap.getOrDefault(num, 0); long sum = sumMap.getOrDefault(num, 0L); arr[i] += sum - (long)count * i; countMap.put(num, count + 1); sumMap.put(num, sum + i); } return arr; }}Complexity
Time
O(N) - The algorithm consists of two independent passes through the array, each taking O(N) time. The total time complexity is O(N) + O(N) = O(N).
Space
O(U) + O(N) - where U is the number of unique elements in `nums`. The maps can store up to U keys. In the worst case, U=N, so the space is O(N). The O(N) term is for the result array.
Trade-offs
Pros
Optimal O(N) time complexity, making it very efficient for large inputs.
The logic is clean, separating the calculation into two distinct and symmetric passes.
Cons
Requires two passes over the input array.
Uses extra space for hash maps, which can be up to O(N) in the worst case.
Solutions
Solution
class Solution {public long[] distance(int[] nums) { int n = nums.length; long[] ans = new long[n]; Map<Integer, List<Integer>> d = new HashMap<>(); for (int i = 0; i < n; ++i) { d.computeIfAbsent(nums[i], k->new ArrayList<>()).add(i); } for (var idx : d.values()) { int m = idx.size(); long left = 0; long right = -1L * m * idx.get(0); for (int i : idx) { right += i; } for (int i = 0; i < m; ++i) { ans[idx.get(i)] = left + right; if (i + 1 < m) { left += (idx.get(i + 1) - idx.get(i)) * (i + 1L); right -= (idx.get(i + 1) - idx.get(i)) * (m - i - 1L); } } } return ans; }}class Solution {public: vector<long long> distance(vector<int> &nums) { int n = nums.size(); vector<long long> ans(n); unordered_map<int, vector<int>> d; for (int i = 0; i < n; ++i) { d[nums[i]].push_back(i); } for (auto &[_, idx] : d) { int m = idx.size(); long long left = 0; long long right = -1LL * m * idx[0]; for (int i : idx) { right += i; } for (int i = 0; i < m; ++i) { ans[idx[i]] = left + right; if (i + 1 < m) { left += (idx[i + 1] - idx[i]) * (i + 1); right -= (idx[i + 1] - idx[i]) * (m - i - 1); } } } return ans; }};class Solution: def distance(self, nums: List[int]) -> List[int]: d = defaultdict(list) for i, x in enumerate(nums): d[x]. append(i) ans = [0] * len(nums) for idx in d . values(): left, right = 0, sum(idx) - len(idx) * idx[0] for i in range(len(idx)): ans[idx[i]] = left + right if i + 1 < len(idx): left += (idx[i + 1] - idx[i]) * (i + 1) right -= (idx[i + 1] - idx[i]) * (len(idx) - i - 1) return ansVideo walkthrough
Newsletter
One sharp idea, every week
System design and interview prep — short enough to finish.
No spam. Unsubscribe anytime.
Practice
Same difficulty — related problems to reinforce the pattern.