Pythonでソート済みリスト内のすべてのペアの絶対差の合計を求めるプログラム

ソートされた数値リスト nums が与えられたとき、リスト内のすべての数値ペアの絶対差の合計を求めることを考えます。ここで、(i, j) と (j, i) は異なるペアとして扱います。答えが非常に大きくなる場合は、結果を 10^9+7 で割った余りを返します。

例えば、nums = [2, 4, 8] の場合、|2 - 4| + |2 - 8| + |4 - 2| + |4 - 8| + |8 - 2| + |8 - 4| を計算することになるため、出力は 24 となります。

解法のアプローチ

この問題を効率的に解くために、以下の手順に従います。

• m = 10^9 + 7 とします。
• total を 0 で初期化します。
• i を 0 から nums のサイズまで繰り返します。
  • total := total + (i * nums[i] - (nums のサイズ - 1 - i) * nums[i]) mod m
• (2 * total) mod m を返します。

考え方のポイント

リストがソートされていることを利用すると、各要素 nums[i] について、それより左側（小さい側）の要素との差は nums[i] - nums[j]、右側（大きい側）の要素との差は nums[j] - nums[i] となります。インデックス i の要素は、左側に i 個、右側に (n - 1 - i) 個の要素を持つため、各要素の合計への寄与は i * nums[i] - (n - 1 - i) * nums[i] と表せます。これにより、全ペアを実際に列挙することなく O(n) の計算量で答えを求められます。最後に、(i, j) と (j, i) の両方を数えるため、合計を 2 倍します。

実装例

以下の実装を見ると、理解がより深まるでしょう。

class Solution:
    def solve(self, nums):
        m = 10**9 + 7
        total = 0
        for i in range(len(nums)):
            total += (i*nums[i] - (len(nums) - 1 - i)*nums[i]) % m
        return (2*total) % m
ob = Solution()
nums = [2, 4, 8]
print(ob.solve(nums))

入力

[2, 4, 8]

出力

24