Problem statement

https://binarysearch.com/problems/Distance-Sums/

Solution

Equal to Leetcode 2121. Intervals Between Identical Elements.

Complexity

It is O(n) for time and space.

Code

class Solution:
    def solve(self, arr):
        pl = defaultdict(list)
        n = len(arr)
        ans = [0] * n
        for i, x in enumerate(arr):
            pl[x] += [i]

        for val in pl:
            a = pl[val]
            t1 = list(accumulate(a))
            t2 = list(accumulate(a[::-1]))[::-1]
            m = len(a)
            b = [0]*m

            for i in range(m):
                if i > 0: b[i] += i * a[i] - t1[i-1]
                if i < m-1: b[i] += t2[i+1]  - (m - i - 1)* a[i]

            for i in range(m):
                ans[a[i]] = b[i]

        return ans