D - 身長差の総和 / Sum of Height Differences Editorial by admin
Claude 4.5 Opus概要
\(N\) 人の生徒から2人を選ぶすべてのペアについて、身長差(絶対値)の総和を効率的に求める問題です。
考察
素朴なアプローチ
最も単純な方法は、すべてのペア \((i, j)\)(\(i < j\))について \(|A_i - A_j|\) を計算して足し合わせることです。
total = 0
for i in range(N):
for j in range(i+1, N):
total += abs(A[i] - A[j])
しかし、ペアの数は \(\frac{N(N-1)}{2}\) 個あり、\(N = 2 \times 10^5\) のとき約 \(2 \times 10^{10}\) 回の計算が必要になります。これでは TLE(時間制限超過) になってしまいます。
重要な気づき:ソートを活用する
配列をソートすると、絶対値を外すことができます。
ソート後の配列を \(A_0 \leq A_1 \leq \ldots \leq A_{N-1}\) とすると、\(i < j\) のとき常に \(A_i \leq A_j\) なので: $\(|A_i - A_j| = A_j - A_i\)$
さらなる効率化:各要素の寄与を考える
ソート後、各要素 \(A_i\) が全体の総和にどれだけ寄与するか考えます。
\(A_i\) は: - 自分より前(小さい) \(i\) 個の要素とペアを作るとき、引かれる側(\(+A_i\) として \(i\) 回寄与) - 自分より後(大きい) \(N-1-i\) 個の要素とペアを作るとき、引く側(\(-A_i\) として \(N-1-i\) 回寄与)
しかし、もう少し簡単に考えましょう。
具体例で理解する
\(A = [3, 1, 4]\) をソートすると \([1, 3, 4]\) になります。
\(i = 2\)(値は \(4\))のとき、前の要素 \([1, 3]\) との差の総和は: $\((4 - 1) + (4 - 3) = 2 \times 4 - (1 + 3) = 8 - 4 = 4\)$
つまり、\(i\) 番目の要素について: $\(i \times A_i - (A_0 + A_1 + \ldots + A_{i-1})\)$
これを累積和を使って効率的に計算できます。
アルゴリズム
- 配列 \(A\) をソートする
- 累積和
prefix_sumを \(0\) で初期化 - 各要素 \(A_i\)(\(i = 0, 1, \ldots, N-1\))について:
- \(A_i\) より前の要素との差の総和 \(= i \times A_i - \text{prefix\_sum}\) を
totalに加算 prefix_sumに \(A_i\) を加算
- \(A_i\) より前の要素との差の総和 \(= i \times A_i - \text{prefix\_sum}\) を
totalを出力
具体例での計算
\(A = [3, 1, 4]\) → ソート後 \([1, 3, 4]\)
| \(i\) | \(A_i\) | \(i \times A_i - \text{prefix\_sum}\) | total | prefix_sum |
|---|---|---|---|---|
| 0 | 1 | \(0 \times 1 - 0 = 0\) | 0 | 1 |
| 1 | 3 | \(1 \times 3 - 1 = 2\) | 2 | 4 |
| 2 | 4 | \(2 \times 4 - 4 = 4\) | 6 | 8 |
答え:\(6\)(検算:\(|3-1| + |4-1| + |4-3| = 2 + 3 + 1 = 6\) ✓)
計算量
- 時間計算量: \(O(N \log N)\)(ソートがボトルネック、その後の処理は \(O(N)\))
- 空間計算量: \(O(N)\)(配列の保持)
実装のポイント
\(A_i\) の最大値が \(10^9\)、\(N\) の最大値が \(2 \times 10^5\) なので、総和は最大で約 \(10^{14}\) 程度になります。Pythonでは整数のオーバーフローを気にする必要はありませんが、C++などでは
long long型を使う必要があります。累積和を別配列で持たず、変数1つで管理することでメモリを節約しています。
ソースコード
def solve():
N = int(input())
A = list(map(int, input().split()))
# ソートする
A.sort()
# 累積和を計算
# ソート後、各要素A[i]について、それより小さい要素との差の総和を効率的に計算
# A[i]とA[j] (j < i)のペアについて、A[i] - A[j]の総和を求める
# = i * A[i] - (A[0] + A[1] + ... + A[i-1])
# = i * A[i] - prefix_sum[i]
total = 0
prefix_sum = 0
for i in range(N):
# A[i]より前の要素との差の総和
# A[i]はi個の要素より大きい(ソート済みなので)
total += i * A[i] - prefix_sum
prefix_sum += A[i]
print(total)
solve()
この解説は claude4.5opus によって生成されました。