Problem Statement in English

You’re given an array of integers nums. You have to return an array res where res[i] is the sum of |i - j| between i and all other j such that nums[i] == nums[j].

Approach

Looking at the constraints we can see that a quadratic solution won’t work. We need something better.

In order to make this efficient we’re going to have to open the mod. This gives us the following:

  • |i - j| = i - j if i > j
  • |i - j| = j - i if i < j

Now we can group the indices of the same values together. Time to do some prefix sums.

To calculate these sums in O(1) time for each i, you pre-calculate the prefix sums of the indices for each unique value.

For a group of k indices, let S be the total sum of all indices in that group. Let P be the prefix sum of indices up to the current index i (including i), and let count be the number of indices processed so far.

The formula for arr[i] becomes: • Left side sum: (count * i) - P • Right side sum: (S - P) - ((k - count) * i)

Final arr[i] = Left side sum + Right side sum

And we’re done!

Solution in Python


class Solution:
    def distance(self, nums: List[int]) -> List[int]:
        hm = defaultdict(list)

        for i, v in enumerate(nums):
            hm[v].append(i)

        res = [0] * len(nums)

        for values in hm.values():
            N = len(values)
            
            if N <= 1:
                continue

            s = sum(values)
            curr = 0
            temp = [0] * N

            for i, index in enumerate(values):
                curr += index

                pre = (index * (i + 1)) - curr
                post = (s - curr) - ((N - (i + 1)) * index)
                
                res[index] = pre + post

        return res

Complexity

  • Time: $O(n)$
    Since we traverse the array a constant number of times.

  • Space: $O(n)$
    Since we store the hashmap and the result array.

Mistakes I Made

I had to look this one up :(

And we are done.