2615. Sum of Distances
Problem Description
In this problem, we're presented with an integer array nums
which is 0-indexed. We're tasked to find another array arr
where each element arr[i]
represents the sum of absolute differences between index i
and every index j
where the value nums[j]
is equal to nums[i]
, excluding when i
and j
are the same. If there are no such j
values (meaning there are no other identical values to nums[i]
in the array), we simply set arr[i]
to 0
.
To illustrate, if nums
was [3, 3, 1]
, then arr
would be [1, 1, 0]
because:
- For
i=0
,nums[0] == 3
, we find another3
at index1
, so the sum is|0 - 1| = 1
. There are no more3
s, soarr[0]
is1
. - For
i=1
, it's the same situation but in reverse, soarr[1]
is also1
. - For
i=2
,nums[2] == 1
and there are no other1
s in the array, soarr[2]
is0
.
We need to perform this calculation efficiently for each index and construct the arr
array accordingly.
Intuition
To find an effective solution, we should try to avoid the naive approach of checking each pair of indices for equality, as that would result in an O(n^2) time complexity, which is impractical for large arrays.
Instead, we can optimize by grouping indices with the same values together. We can use a dictionary to map each unique number in the array to the list of indices where that number appears. This will help us compute the sum of distances for each group in aggregate, rather than individually.
Once we have those index lists, we can iterate over them to calculate the sum of distances for each index in the group. Here's the step-by-step approach:
- We compute the sum of indices for each group - this gives us a starting sum for the right part of each index.
- We start from the leftmost index in each group and iterate towards the right.
- For each index, we maintain two variables:
left
andright
, which represent the sum of distances to indices on the left and right, respectively. - As we move from one index to the next within a group, we update
left
by adding the distance to the next index multiplied by the count of indices that are now on the left. - We update
right
by subtracting the same distance multiplied by the count of indices remaining on the right. - We use these
left
andright
accumulators to calculate the total distance for each index, this total distance is what we add to the corresponding position in thearr
array.
This approach reduces the amount of redundant calculations and brings the complexity down significantly because each distance is essentially calculated in linear time relative to the number of occurrences of each unique value.
Learn more about Prefix Sum patterns.
Solution Approach
The solution utilizes a hash map (via Python's defaultdict
from the collections
module) to group the indices of identical values in nums
. Here's how the solution is implemented:
-
Initialize a
defaultdict
of lists namedd
. Iterate overnums
usingenumerate
to build the dictionary where each key is a value innums
, and the corresponding value is a list of indices where that key appears. -
Create an array
ans
filled with zeros of the same length asnums
. This array will store our resulting distances for each index. -
For each list of indices
idx
in the hash map:- Initialize
left
as0
because there are no indices to the left of our starting point. - Calculate the initial value of
right
by summing up all the indices inidx
and adjusting for the lowest index (len(idx) * idx[0]
). This variable represents the sum of distances of all indices to the right of our current index.
- Initialize
-
Iterate over each index
i
inidx
:- Set
ans[idx[i]]
to be the sum ofleft
plusright
, which provides the sum of distances fornums[idx[i]]
. - Before we move to the next index
i+1
, we updateleft
andright
. Updateleft
by adding the distance to the next index multiplied byi + 1
(because we'll have one more element to the left). - Update
right
by subtracting the same distance multiplied by(len(idx) - i - 1)
to account for having one less element to the right.
- Set
-
Loop until all indices in
idx
are processed. By the end,ans
will have the calculated distances for each element innums
.
This approach is efficient because it computes the distances for indices of the same value in nums
as a whole, minimizing repeated work. By using left
and right
accumulators, it avoids iterating over all pairs of indices and scales linearly with the number of occurrences of each unique value in nums
.
The algorithmic patterns used here include hashing (to group indices), enumeration (to iterate with indices), and prefix sums (though not explicitly stored, they are implied in the use of left
and right
to keep track of cumulative distances).
Ready to land your dream job?
Unlock your dream job with a 2-minute evaluator for a personalized learning plan!
Start EvaluatorExample Walkthrough
For illustration, let's use nums = [2, 1, 2, 1, 2]
to demonstrate the solution approach.
-
We first create the hash map with the list of indices for each value in
nums
:For value
2
, the list of indices is[0, 2, 4]
. For value1
, the list of indices is[1, 3]
.Our hash map (
d
) now looks like this:{ 2: [0, 2, 4], 1: [1, 3] }
-
We then initialize the
ans
array to[0, 0, 0, 0, 0]
, corresponding to the five elements innums
. -
We begin processing the list of indices for the value
2
:-
Initial setup for value
2
:left = 0
,right
is calculated by(2 + 4) - (2 * 0) = 6
. -
First index (i=0):
- We set
ans[0] = left + right
→ans[0] = 0 + 6 = 6
. - We update
left
to be0 + (1 * 2)
since we're moving one step to the right, adding two spaces (distance) from the first index. - We update
right
by subtracting2 * (3 - 1 - 0)
to account for the two spaces we're moving away from the remaining two indices (2 and 4).
- We set
-
Second index (i=1):
- Now
left = 2
andright = 2
(from previous calculations). - Thus,
ans[2] = left + right
→ans[2] = 2 + 2 = 4
. - Update
left
to be2 + (2 * 2)
as we're now considering the total distance from the first and second index to the third. - Update
right
to be2 - (2 * (2 - 1 - 1))
, no change because only one index remains to the right, and we're moving one step away from it.
- Now
-
Third index (i=2):
- Now
left = 6
andright = 0
(there are no more indices to the right). - So,
ans[4] = left + right
→ans[4] = 6 + 0 = 6
.
- Now
With value
2
processed,ans
becomes[6, 0, 4, 0, 6]
. -
-
Next, proceed with the indices for value
1
:-
Initial setup for value
1
:left = 0
,right
is calculated by(3) - (1 * 1) = 2
. -
First index (i=0):
ans[1] = left + right
→ans[1] = 0 + 2 = 2
.- Update
left
to be0 + (1 * 1)
, - Update
right
to be2 - (1 * (2 - 1 - 0))
, which makesright = 1
now.
-
Second index (i=1):
left = 1
andright = 0
.ans[3] = left + right
→ans[3] = 1 + 0 = 1
.
Now, our answer array
ans
is completely updated to[6, 2, 4, 1, 6]
. -
After processing both lists of indices, we get arr = [6, 2, 4, 1, 6]
corresponding to the sum of absolute differences for each index in nums
, as intended by the solution approach.
Solution Implementation
1from collections import defaultdict
2from typing import List
3
4class Solution:
5 def distance(self, nums: List[int]) -> List[int]:
6 # Create a dictionary to store indices of each number in nums
7 indices_dict = defaultdict(list)
8 for index, num in enumerate(nums):
9 indices_dict[num].append(index)
10
11 # Initialize the answer list with zeros
12 distances = [0] * len(nums)
13
14 # Iterate over the groups of indices for each unique number
15 for indices in indices_dict.values():
16 # Initialize the distance from the left and right side
17 left_distance = 0
18 right_distance = sum(indices) - len(indices) * indices[0]
19
20 # Iterate through the indices to compute distances
21 for idx_position in range(len(indices)):
22 # The total distance is the sum of left and right distances
23 distances[indices[idx_position]] = left_distance + right_distance
24
25 # Update the left and right distances for the next index (if not at the end)
26 if idx_position + 1 < len(indices):
27 gap = indices[idx_position + 1] - indices[idx_position]
28 left_distance += gap * (idx_position + 1)
29 right_distance -= gap * (len(indices) - idx_position - 1)
30
31 # Return the list of distances
32 return distances
33
1class Solution {
2 public long[] distance(int[] nums) {
3 int n = nums.length; // Get the length of the input array.
4 long[] answer = new long[n]; // Initialize the answer array to store the final distances.
5 Map<Integer, List<Integer>> indexMap = new HashMap<>(); // Map to store the indices of each number in nums.
6
7 // Loop through nums to fill the indexMap.
8 for (int i = 0; i < n; ++i) {
9 // If the key doesn't exist, create a new list, then add the index.
10 indexMap.computeIfAbsent(nums[i], k -> new ArrayList<>()).add(i);
11 }
12
13 // Iterate over the entry set of the map to calculate distances.
14 for (List<Integer> indices : indexMap.values()) {
15 int m = indices.size(); // The frequency of the current number.
16 long leftSum = 0; // Sum to store the left side distances.
17 long rightSum = -1L * m * indices.get(0); // Sum to store the right side distances, initially inverted.
18
19 // Initial rightSum calculation by summing all indices.
20 for (int index : indices) {
21 rightSum += index;
22 }
23
24 // Iterate over each index for the current number.
25 for (int i = 0; i < m; ++i) {
26 // Calculate the final distance for the current index.
27 int currentIndex = indices.get(i);
28 answer[currentIndex] = leftSum + rightSum;
29
30 // Recalculate leftSum and rightSum for the next index if it exists.
31 if (i + 1 < m) {
32 int nextIndex = indices.get(i + 1);
33 long diff = nextIndex - currentIndex; // Difference between the current and the next index.
34 leftSum += diff * (i + 1L);
35 rightSum -= diff * (m - i - 1L);
36 }
37 }
38 }
39
40 return answer; // Return the completed answer array.
41 }
42}
43
1class Solution {
2public:
3 // Function to calculate distances
4 vector<long long> distance(vector<int>& nums) {
5 // n holds the length of the input array nums
6 int n = nums.size();
7 // Create a vector to store the final answer
8 vector<long long> ans(n);
9 // Use a map to store indices of elements in nums
10 unordered_map<int, vector<int>> indexMap;
11
12 // Populate the map with the indices of each element
13 for (int i = 0; i < n; ++i) {
14 indexMap[nums[i]].push_back(i);
15 }
16
17 // Iterate over each element of the map
18 for (auto& [value, indices] : indexMap) {
19 // m holds the number of occurrences of a particular value in nums
20 int m = indices.size();
21 // left is used to keep track of the cumulative distance from the left
22 long long left = 0;
23 // right is used to keep track of the cumulative distance from the right
24 long long right = -1LL * m * indices[0];
25 // Calculate the initial right value by summing up all indices
26 for (int index : indices) {
27 right += index;
28 }
29 // Iterate over all occurrences of the current value
30 for (int i = 0; i < m; ++i) {
31 // Store the total distance at the current index
32 ans[indices[i]] = left + right;
33
34 // Update left and right for the next values if not at the end
35 if (i + 1 < m) {
36 left += (indices[i + 1] - indices[i]) * (i + 1);
37 right -= (indices[i + 1] - indices[i]) * (m - i - 1);
38 }
39 }
40 }
41 // Return the final vector of distances
42 return ans;
43 }
44};
45
1// Function to calculate distances
2function distance(nums: number[]): number[] {
3 // n holds the length of the input array nums
4 let n: number = nums.length;
5 // Create an array to store the final answer
6 let ans: number[] = new Array(n);
7 // Use a map to store indices of elements in nums
8 let indexMap: Map<number, number[]> = new Map<number, number[]>();
9
10 // Populate the map with the indices of each element
11 for (let i = 0; i < n; ++i) {
12 if (!indexMap.has(nums[i])) {
13 indexMap.set(nums[i], []);
14 }
15 indexMap.get(nums[i])!.push(i);
16 }
17
18 // Iterate over each element of the map
19 indexMap.forEach((indices: number[], value: number) => {
20 // m holds the number of occurrences of a particular value in nums
21 let m: number = indices.length;
22 // left is used to keep track of the cumulative distance from the left
23 let left: number = 0;
24 // right is used to keep track of the cumulative distance from the right
25 let right: number = -1 * m * indices[0];
26 // Calculate the initial right value by summing up all indices
27 for (let index of indices) {
28 right += index;
29 }
30 // Iterate over all occurrences of the current value
31 for (let i = 0; i < m; ++i) {
32 // Store the total distance at the current index
33 ans[indices[i]] = left + right;
34
35 // Update left and right for the next values if not at the end
36 if (i + 1 < m) {
37 left += (indices[i + 1] - indices[i]) * (i + 1);
38 right -= (indices[i + 1] - indices[i]) * (m - i - 1);
39 }
40 }
41 });
42
43 // Return the final array of distances
44 return ans;
45}
46
Time and Space Complexity
The above code maintains a dictionary to store indices of each unique element in the list nums
and then computes the distance sum for each unique element.
Time Complexity:
- The first for loop iterates over all
n
elements of the listnums
to create the dictionaryd
, which takesO(n)
time. - The second loop processes each group of indices stored in the dictionary
d
. Letk
be the maximum size of these index groups. For each group, it iterates over all elements in the list of indices which, in the worst case, could takeO(k)
. Since there are as many groups as there are unique elements innums
, we can say that it takesO(u * k)
time, whereu
is the number of unique elements. - Inside this loop, it performs constant-time operations like addition and subtraction for updating left and right variables.
Since k
times the number of unique elements u
cannot exceed n
(as each element is part of one group and is visited once), the second loop's time complexity is bounded by O(n)
.
Therefore, the total time complexity is O(n)
+ O(n)
= O(n)
.
Space Complexity:
- The space taken by the dictionary
d
depends on the number of unique elements innums
. If there areu
unique elements, it could take up toO(u * k)
space, wherek
is the average size of the index lists for each unique element. However, since there aren
total elements,u * k
is also bounded byn
. Therefore, it usesO(n)
space. - We also have the
ans
array, which is equal in size tonums
, adding anotherO(n)
space.
Adding these together, the space complexity of the code is O(n)
+ O(n)
= O(n)
.
Learn more about how to find time and space complexity quickly using problem constraints.
How would you design a stack which has a function min
that returns the minimum element in the stack, in addition to push
and pop
? All push
, pop
, min
should have running time O(1)
.
Recommended Readings
Prefix Sum The prefix sum is an incredibly powerful and straightforward technique Its primary goal is to allow for constant time range sum queries on an array What is Prefix Sum The prefix sum of an array at index i is the sum of all numbers from index 0 to i By
LeetCode Patterns Your Personal Dijkstra's Algorithm to Landing Your Dream Job The goal of AlgoMonster is to help you get a job in the shortest amount of time possible in a data driven way We compiled datasets of tech interview problems and broke them down by patterns This way we
Recursion Recursion is one of the most important concepts in computer science Simply speaking recursion is the process of a function calling itself Using a real life analogy imagine a scenario where you invite your friends to lunch https algomonster s3 us east 2 amazonaws com recursion jpg You first
Want a Structured Path to Master System Design Too? Don’t Miss This!