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 another 3 at index 1, so the sum is |0 - 1| = 1. There are no more 3s, so arr[0] is 1.
  • For i=1, it's the same situation but in reverse, so arr[1] is also 1.
  • For i=2, nums[2] == 1 and there are no other 1s in the array, so arr[2] is 0.

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:

  1. We compute the sum of indices for each group - this gives us a starting sum for the right part of each index.
  2. We start from the leftmost index in each group and iterate towards the right.
  3. For each index, we maintain two variables: left and right, which represent the sum of distances to indices on the left and right, respectively.
  4. 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.
  5. We update right by subtracting the same distance multiplied by the count of indices remaining on the right.
  6. We use these left and right accumulators to calculate the total distance for each index, this total distance is what we add to the corresponding position in the arr 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:

  1. Initialize a defaultdict of lists named d. Iterate over nums using enumerate to build the dictionary where each key is a value in nums, and the corresponding value is a list of indices where that key appears.

  2. Create an array ans filled with zeros of the same length as nums. This array will store our resulting distances for each index.

  3. For each list of indices idx in the hash map:

    • Initialize left as 0 because there are no indices to the left of our starting point.
    • Calculate the initial value of right by summing up all the indices in idx 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.
  4. Iterate over each index i in idx:

    • Set ans[idx[i]] to be the sum of left plus right, which provides the sum of distances for nums[idx[i]].
    • Before we move to the next index i+1, we update left and right. Update left by adding the distance to the next index multiplied by i + 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.
  5. Loop until all indices in idx are processed. By the end, ans will have the calculated distances for each element in nums.

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 Evaluator

Example Walkthrough

For illustration, let's use nums = [2, 1, 2, 1, 2] to demonstrate the solution approach.

  1. 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 value 1, the list of indices is [1, 3].

    Our hash map (d) now looks like this:

    {
        2: [0, 2, 4],
        1: [1, 3]
    }
  2. We then initialize the ans array to [0, 0, 0, 0, 0], corresponding to the five elements in nums.

  3. 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 + rightans[0] = 0 + 6 = 6.
      • We update left to be 0 + (1 * 2) since we're moving one step to the right, adding two spaces (distance) from the first index.
      • We update right by subtracting 2 * (3 - 1 - 0) to account for the two spaces we're moving away from the remaining two indices (2 and 4).
    • Second index (i=1):

      • Now left = 2 and right = 2 (from previous calculations).
      • Thus, ans[2] = left + rightans[2] = 2 + 2 = 4.
      • Update left to be 2 + (2 * 2) as we're now considering the total distance from the first and second index to the third.
      • Update right to be 2 - (2 * (2 - 1 - 1)), no change because only one index remains to the right, and we're moving one step away from it.
    • Third index (i=2):

      • Now left = 6 and right = 0 (there are no more indices to the right).
      • So, ans[4] = left + rightans[4] = 6 + 0 = 6.

    With value 2 processed, ans becomes [6, 0, 4, 0, 6].

  4. 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 + rightans[1] = 0 + 2 = 2.
      • Update left to be 0 + (1 * 1),
      • Update right to be 2 - (1 * (2 - 1 - 0)), which makes right = 1 now.
    • Second index (i=1):

      • left = 1 and right = 0.
      • ans[3] = left + rightans[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 list nums to create the dictionary d, which takes O(n) time.
  • The second loop processes each group of indices stored in the dictionary d. Let k 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 take O(k). Since there are as many groups as there are unique elements in nums, we can say that it takes O(u * k) time, where u 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 in nums. If there are u unique elements, it could take up to O(u * k) space, where k is the average size of the index lists for each unique element. However, since there are n total elements, u * k is also bounded by n. Therefore, it uses O(n) space.
  • We also have the ans array, which is equal in size to nums, adding another O(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.


Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:
Question 1 out of 10

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

Want a Structured Path to Master System Design Too? Don’t Miss This!