2615. Sum of Distances

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.

 

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

 

Note: This question is the same as 2121: Intervals Between Identical Elements.

 

Solutions

Solution: Prefix Sum

• Time complexity: O(n)
• Space complexity: O(n)

 

JavaScript

/**
 * @param {number[]} nums
 * @return {number[]}
 */
const distance = function (nums) {
  const n = nums.length;
  const prefixMap = new Map();
  const suffixMap = new Map();
  const result = Array.from({ length: n }, () => 0);

  for (let index = 0; index < n; index++) {
    const num = nums[index];

    if (prefixMap.has(num)) {
      const { sum, count } = prefixMap.get(num);

      result[index] += index * count - sum;
      prefixMap.set(num, { sum: sum + index, count: count + 1 });
    } else {
      prefixMap.set(num, { sum: index, count: 1 });
    }
  }

  for (let index = n - 1; index >= 0; index--) {
    const num = nums[index];

    if (suffixMap.has(num)) {
      const { sum, count } = suffixMap.get(num);

      result[index] += Math.abs(index * count - sum);
      suffixMap.set(num, { sum: sum + index, count: count + 1 });
    } else {
      suffixMap.set(num, { sum: index, count: 1 });
    }
  }

  return result;
};