2615

2615. Sum of Distances (Medium)

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 <= 105
• 0 <= nums[i] <= 109

Companies: BNY Mellon

Related Topics:
Array, Hash Table, Prefix Sum

Similar Questions:

• Remove Duplicates from Sorted Array (Easy)
• Find All Duplicates in an Array (Easy)
• Minimum Operations to Make All Array Elements Equal (Easy)

Solution 1. Left-to-right State Transition

Group the indices of the same numbers together using an unordered_map<int, vector<int>> m.

For the index numbers in each of the group, we keep rightSum as sum of all index numbers greater than the current index number, and leftSum as the sum of all index numbers smaller than or equal to the current index number. Initially rightSum is the sum of all index numbers in the group.

Then, we visit the index numbers one by one. For index number v[i], we deduct it from rightSum and add it to leftSum. ans[v[i]] is rightSum - (N - i - 1) * v[i] + (i + 1) * v[i] - leftSum = rightSum - leftSum + (2 * (i + 1) - N) * v[i].

// OJ: https://leetcode.com/problems/sum-of-distances
// Author: github.com/lzl124631x
// Time: O(N)
// Space: O(N)
class Solution {
public:
    vector<long long> distance(vector<int>& A) {
        int N = A.size();
        vector<long long> ans(N);
        unordered_map<int, vector<int>> m;
        for (int i = 0; i < N; ++i) m[A[i]].push_back(i);
        for (auto &[n, v] : m) {
            long long rightSum = 0, leftSum = 0;
            for (int i : v) rightSum += i;
            for (int i = 0; i < v.size(); ++i) {
                rightSum -= v[i];
                leftSum += v[i];
                ans[v[i]] = rightSum - leftSum + (2LL * (i + 1) - v.size()) * v[i];
            }
        }
        return ans;
    }
};