Problem statement

https://leetcode.com/problems/intervals-between-identical-elements/

Solution

For each value find all places where we have this value. Then use cumulative sums to find sums.

Complexity

It is O(n) for time and space.

Code

class Solution:
    def getDistances(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