Sum of Distances

Problem Description

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] 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 <= 10^5
• 0 <= nums[i] <= 10^9

Approach

• lets take a example [1,3,1,1,2,1,1]
• lets take value 1 the arr formed for value 1 will be:[0,2,3,5,6]
• lets calculate ans for index 3:
• it will be divided into two part one which is on left and other on right
• for left: val1=(3-0)+(3-2) = (3+3)-(0+2) = 3*left_size-(sum upto i-1);
• for right: val2=(6-3)+(5-3) = (6+5)-3*2 = (sum from i+1)-(right_size * 3)
• now ans for 3 index=val1+val2;

Complexity

Time complexity: $O(n*logn)$ Space complexity: $O(n)$

Solution

Code in Different Languages

C++

class Solution {
public:
   vector<long long> distance(vector<int>& nums) {
       unordered_map<int,vector<int>>mp;
       for(int i=0;i<nums.size();i++)
           mp[nums[i]].push_back(i);
       
       
       unordered_map<int,vector<long long>>dp;//for storing prefix sum
       
       for(auto it:mp)
       {
           auto arr=it.second;
           dp[it.first].push_back(arr[0]);
           for(int i=1;i<arr.size();i++)
           {
               dp[it.first].push_back(0ll+dp[it.first].back()+arr[i]);
           }
       }
       
       vector<long long>ans(nums.size(),0);
       
       
       for(auto it:mp)
       {
           auto arr1=it.second;
           auto arr2=dp[it.first];
           
           for(int i=arr1.size()-1;i>=0;i--)
           {
               int left=i;
               int right=arr1.size()-i-1;
               
               long long sum1=i-1>=0?arr2[i-1]:0;
               long long sum2=arr2[arr1.size()-1]-arr2[i];
               
               long long val=(1ll*left*arr1[i]-sum1)+(sum2-1ll*right*arr1[i]);
               ans[arr1[i]]=val;
           }
       }
       return ans;
   }
};

JAVA

class Solution {
    public long[] distance(int[] nums) {
        HashMap<Integer, ArrayList<Integer>> hm = new HashMap<>();
        for(int i = 0; i < nums.length; i++) {
            if(!hm.containsKey(nums[i])) {
                hm.put(nums[i], new ArrayList<>());
            }
            hm.get(nums[i]).add(i);
        }
        
        long[] ans = new long[nums.length];
        for(ArrayList<Integer> indices : hm.values()) {
            int n = indices.size();
            long sum = 0;
            for(int i = 0; i < n; i++) {
                sum += indices.get(i);
            }
            long leftsum = 0;
            long rightsum = sum;
            long currsum = 0;
            for(int i = 0; i < n; i++) {
                currsum = 0;
                currsum += (long)i*(long)indices.get(i) - leftsum;
                currsum += rightsum - (long)(n-i)*(long)indices.get(i);
                leftsum += (long)indices.get(i);
                rightsum -= (long)indices.get(i);
                ans[indices.get(i)] = currsum;
            }
        }
        return ans;
    }
}

PYTHON

class Solution:
    def distance(self, nums: List[int]) -> List[int]:
        num_indices = dict()
        occ=dict()
        for i, num in enumerate(nums):
            if num not in num_indices:
                num_indices[num] = i
                occ[num]=1
            else:
                num_indices[num]=num_indices[num]+i
                occ[num]=occ[num]+1
        arr = [0] * len(nums)
        n=len(nums)
        for i in range(n):
            arr[i] = num_indices[nums[i]] - occ[nums[i]]*i     
            num_indices[nums[i]]=num_indices[nums[i]]-2*i
            occ[nums[i]]=occ[nums[i]]-2
            
        return arr

Complexity Analysis

• Time Complexity: $O(n*logn)$

• Space Complexity: $O(n)$

References

• LeetCode Problem: Sum of Distances