2615. Sum of Distances - LeetCode

from collections import defaultdict
from itertools import accumulate


class Solution0:
    def distance(self, nums: list[int]) -> list[int]:
        groups = defaultdict(list[int])
        for pos, num in enumerate(nums):
            groups[num].append(pos)

        def prefix_sum(arr: list[int]) -> list[int]:
            result = [0]
            for x in arr:
                result.append(result[-1] + x)
            return result

        def distince(
            rank: int,
            pos: int,
            prefix: list[int],
        ) -> int:
            # start ~ rank
            a, b = 0, rank
            left = pos * (b - a) - (prefix[b] - prefix[a])

            # next (rank+1) ~ last
            a, b = rank + 1, len(prefix) - 1
            right = -pos * (b - a) + (prefix[b] - prefix[a])

            return left + right

        result = [0] * len(nums)
        for pos_arr in groups.values():
            prefix = prefix_sum(pos_arr)
            for rank, pos in enumerate(pos_arr):
                result[pos] = distince(
                    rank,
                    pos,
                    prefix,
                )
        return result


class Solution1:
    def distance(self, nums: list[int]) -> list[int]:
        groups = defaultdict(list[int])
        for pos, num in enumerate(nums):
            groups[num].append(pos)

        result = [0] * len(nums)
        for pos_arr in groups.values():
            prefix = list(accumulate(pos_arr))
            last_i = len(pos_arr) - 1
            for rank, pos in enumerate(pos_arr):
                # prefix: 1   [4]   11     21    33
                #
                # rank:   0    1    [2]     3     4
                # pos:    1    3     7     10    12
                #        7-1  7-3    ^    10-7  12-7
                #
                #           7*2         -7*((len-1)-2)
                #     -prefix[2]+7     +prefix[(len-1)-2]
                #

                left = pos * (rank) - (prefix[rank]) + pos
                right = -pos * (last_i - rank) + (prefix[last_i] - prefix[rank])
                result[pos] = left + right
        return result


class Solution:
    def distance(self, nums: list[int]) -> list[int]:
        groups = defaultdict(list[int])
        for pos, num in enumerate(nums):
            groups[num].append(pos)

        result = [0] * len(nums)
        for pos_arr in groups.values():
            prefix = list(accumulate(pos_arr))
            last_i = len(pos_arr) - 1
            for rank, pos in enumerate(pos_arr):
                result[pos] = pos * (2 * rank - last_i + 1) + prefix[last_i] - 2 * prefix[rank]
        return result