Official

D - 身長差の総和 / Sum of Height Differences Editorial by admin Editorial for printing

GPT 5.2 High

概要

与えられた \(N\) 人の身長 \(A_i\) について、すべての組 \((i<j)\) の身長差 \(|A_i-A_j|\) の総和を高速に求めます。

考察

素朴に全ペアを調べると、組の数は \(\frac{N(N-1)}{2}\) 個あり、計算量は \(O(N^2)\) になります。
\(N \le 2\times 10^5\) なので、\(N^2\) は最大で約 \(4\times 10^{10}\) 回の計算になり、時間内に終わりません(TLE)。

ここで重要な気づきは次の2つです。

  1. 身長をソートすると絶対値が外れる
    身長を昇順に並べて \(a_0 \le a_1 \le \dots \le a_{N-1}\) とすると、\(i<j\) のとき
    \(|a_j-a_i| = a_j-a_i\) となり、絶対値を考えずに差をそのまま足せます。

  2. 各要素の「貢献度」をまとめて計算できる
    固定した \(a_i\) について、左側(小さい値)にある要素との身長差の総和は [ \sum_{k=0}^{i-1} (a_i - ak) ] です。これは [ \sum{k=0}^{i-1} (a_i - a_k) = i\cdot ai - \sum{k=0}^{i-1} ak ] と変形でき、左側の総和 $\sum{k=0}^{i-1} a_k\((累積和)が分かれば \)O(1)$ で求まります。

具体例

例えば \([160, 150, 170]\) をソートして \([150,160,170]\) とします。

  • \(i=0\)(150): 左がないので加算なし
  • \(i=1\)(160): \((160-150)=10\)
  • \(i=2\)(170): \((170-150)+(170-160)=30\)

合計 \(10+30=40\) となり、全ペアの絶対差の総和と一致します。

アルゴリズム

  1. 配列 \(A\) を昇順にソートする。
  2. 変数 pref を「これまでに見た要素の総和(累積和)」として持つ。
  3. 左から順に \(i=0..N-1\) について次を行う:
    • 現在の値を \(x=a_i\) とすると、左側との身長差総和は
      [ x\cdot i - \text{pref} ] なので、これを答え ans に加える。
    • pref += x として累積和を更新する。
  4. ans を出力する。

この方法で、各要素につき \(O(1)\) で寄与を計算でき、全体を高速化できます。

計算量

  • 時間計算量: \(O(N\log N)\)(ソートが支配的。走査部分は \(O(N)\))
  • 空間計算量: \(O(N)\)(配列を保持)

実装のポイント

  • ソート後は \(i<j\) なら差が必ず非負なので、絶対値を外して扱えます。

  • ans は最大で大きな値になる可能性があります(\(A_i\) は最大 \(10^9\)、組数も多い)。Python の int は任意精度なのでそのままでOKです(他言語なら 64bit 整数が必要)。

  • 入力が大きいので、sys.stdin.buffer.read() を使うと高速に読み込めます。

    ソースコード

import sys

def main():
    data = sys.stdin.buffer.read().split()
    n = int(data[0])
    a = list(map(int, data[1:1+n]))
    a.sort()

    pref = 0
    ans = 0
    for i, x in enumerate(a):
        ans += x * i - pref
        pref += x

    print(ans)

if __name__ == "__main__":
    main()

この解説は gpt-5.2-high によって生成されました。

posted:
last update: