Skip to main content

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)O(n*logn) Space complexity: O(n)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)O(n*logn)

  • Space Complexity: O(n)O(n)

References​

  • LeetCode Problem: Sum of Distances