Editorial

Solution walkthrough

Review the explanation, then compare the reference implementation.

Video explanation

Solution

Intuition

For equal values encountered in index order, moving from their previous occurrence to the current one increases the distance to every earlier occurrence by the same gap. Keep the previous distance sum and multiply that gap by the number of earlier equal values. A reverse pass handles later occurrences symmetrically.

Approach

In the forward pass, store per value its occurrence count, last index, and accumulated distance at that index. Set the current prefix distance to the old sum plus count * indexGap, then update the record. Reset the map and repeat right to left for suffix distances. Add the two arrays position by position. Values occurring once contribute zero, and other values' positions never enter a group's calculation.

class Solution:
    def distance(self, nums: List[int]) -> List[int]:
        n = len(nums)
        mp = {}
        # n -> [cnt, last_i, acc_d]
        pre = [0] * n
        suf = [0] * n

        for i, x in enumerate(nums):
            if x in mp:
                pre[i] = mp[x][2] + (mp[x][0] * (i - mp[x][1]))
                mp[x][0] += 1
                mp[x][1] = i
                mp[x][2] = pre[i]
            else:
                mp[x] = [1, i, 0]
                pre[i] = 0

        mp = {}
        for i in range(n - 1, -1, -1):
            x = nums[i]
            if x in mp:
                suf[i] = mp[x][2] + (mp[x][0] * (mp[x][1] - i))
                mp[x][0] += 1
                mp[x][1] = i
                mp[x][2] = suf[i]
            else:
                mp[x] = [1, i, 0]
                suf[i] = 0

        return [pre[i] + suf[i] for i in range(n)]

Complexity analysis

Let n = nums.length and u be the number of distinct values. Time: O(n) across the forward and backward grouped-distance scans. Auxiliary space: O(n + u) for prefix/suffix arrays and per-value state; the returned array uses O(n) space.