D - 都市間の距離 / Distance Between Cities Editorial by admin
Claude 4.5 Opus概要
\(N\) 個の都市間のマンハッタン距離の総和を、各次元ごとに独立に計算し、ソートと累積和を用いて効率的に求める問題です。
考察
素朴なアプローチとその問題点
まず思いつくのは、すべての都市ペア \((i, j)\) について愚直にマンハッタン距離を計算する方法です。
for i in range(N):
for j in range(i+1, N):
for k in range(M):
total += abs(A[i][k] - A[j][k])
この方法では、都市ペアの数が \(\frac{N(N-1)}{2}\) 個あり、各ペアで \(M\) 次元の計算が必要なので、計算量は \(O(N^2 M)\) となります。\(N = 2 \times 10^5\) の場合、これは約 \(2 \times 10^{11}\) 回の演算となり、TLE(時間制限超過)になってしまいます。
重要な気づき:次元ごとに分離できる
マンハッタン距離の定義をよく見ると、各次元の絶対値の差の和になっています。つまり、求める値は次のように変形できます:
\[\sum_{1 \leq i < j \leq N} \sum_{k=1}^{M} |A_{i,k} - A_{j,k}| = \sum_{k=1}^{M} \sum_{1 \leq i < j \leq N} |A_{i,k} - A_{j,k}|\]
各次元 \(k\) ごとに独立に計算できるのです!
1次元での効率的な計算
次元 \(k\) について、座標値を \(x_1, x_2, \ldots, x_N\) とし、これをソートして \(y_1 \leq y_2 \leq \cdots \leq y_N\) とします。
ソート後は \(|y_i - y_j| = y_i - y_j\)(\(i > j\) のとき)と絶対値を外せます。
\(y_i\) が総和に寄与する分を考えると: - \(y_i\) より小さい要素は \(i-1\) 個ある - それぞれとのペアで \(y_i - y_j\) を足す - よって \(y_i\) の寄与は \((i-1) \times y_i - (y_1 + y_2 + \cdots + y_{i-1})\)
累積和を使えば、これを \(O(1)\) で計算できます。
具体例
座標が \([3, 1, 4]\) の場合: 1. ソート後:\([1, 3, 4]\) 2. \(i=0\):寄与 = \(0 \times 1 - 0 = 0\) 3. \(i=1\):寄与 = \(1 \times 3 - 1 = 2\)(ペア \((1,3)\) の差) 4. \(i=2\):寄与 = \(2 \times 4 - 4 = 4\)(ペア \((1,4)\) と \((3,4)\) の差の和)
総和 = \(0 + 2 + 4 = 6 = |3-1| + |3-4| + |1-4|\) ✓
アルゴリズム
- 各次元 \(k = 1, 2, \ldots, M\) について以下を行う:
- 全都市の第 \(k\) 次元の座標を配列に集める
- その配列をソートする
- 累積和を更新しながら、各要素について「自分より小さい全要素との差の和」を計算
- 総和に加算
- 最終的な総和を出力する
計算量
- 時間計算量: \(O(MN \log N)\)
- 各次元でソートに \(O(N \log N)\)、累積和の計算に \(O(N)\)
- これを \(M\) 次元分繰り返す
- 空間計算量: \(O(N)\)
- 各次元の座標を格納する配列
実装のポイント
オーバーフロー対策:座標の絶対値が最大 \(10^6\)、都市数が最大 \(2 \times 10^5\) なので、総和は非常に大きくなりえます。Python では整数のオーバーフローがないので安心ですが、他の言語では64ビット整数型を使用する必要があります。
累積和の更新タイミング:\(i\) 番目の要素の寄与を計算した後に累積和を更新することで、「自分より前の要素の和」を正しく取得できます。
ソースコード
def solve():
N, M = map(int, input().split())
A = []
for _ in range(N):
row = list(map(int, input().split()))
A.append(row)
total = 0
# 各次元ごとに独立に計算
for j in range(M):
# j番目の座標成分を取り出してソート
coords = [A[i][j] for i in range(N)]
coords.sort()
# ソート後、i番目の要素は、それより前のi個の要素との差の絶対値の和に寄与
# coords[i] - coords[k] (k < i) の総和を効率的に計算
# = i * coords[i] - (coords[0] + coords[1] + ... + coords[i-1])
prefix_sum = 0
for i in range(N):
# coords[i]は、前のi個の要素それぞれとのペアで、coords[i] - coords[k]を寄与
total += i * coords[i] - prefix_sum
prefix_sum += coords[i]
print(total)
solve()
この解説は claude4.5opus によって生成されました。
posted:
last update: