D - 都市間の距離 / Distance Between Cities 解説 by admin
GPT 5.2 High概要
\(N\) 個の都市の \(M\) 次元座標が与えられたとき、すべての組 \((i<j)\) についてマンハッタン距離 \(\sum_{k=1}^{M}|A_{i,k}-A_{j,k}|\) を足し合わせた総和を高速に求めます。
考察
マンハッタン距離は次元ごとの絶対値差の和なので、総和も次元ごとに分解できます。
\( \sum_{1 \le i < j \le N}\sum_{k=1}^{M}|A_{i,k}-A_{j,k}| = \sum_{k=1}^{M}\sum_{1 \le i < j \le N}|A_{i,k}-A_{j,k}| \)
つまり「\(M\) 回、1次元の問題(全ペアの \(|x_i-x_j|\) の総和)を解いて足す」だけで良いです。
しかし素朴に全ペアを列挙すると、組の数は \(\binom{N}{2}\approx 2\times 10^5 \times 10^5\) で非常に大きく、 \( O(N^2M) \) となって確実に間に合いません(TLE)。
そこで1次元について、ソートと累積和で \( \sum_{i<j}|x_i-x_j| \) を \(O(N\log N)\) で計算します。
1次元での重要な気づき
数列 \(x\) を昇順にソートして \(a_0 \le a_1 \le \dots \le a_{N-1}\) とすると、\(i<j\) では常に \(|a_j-a_i|=a_j-a_i\) です。
よって、各 \(a_i\) が「右側の要素との差」に何回登場するかを考えると、
- \(a_i\) は左側の要素(\(i\) 個)との差で「大きい側」として出てくる
→ 寄与は \(\sum_{t=0}^{i-1} (a_i-a_t)= i\cdot a_i - \sum_{t=0}^{i-1}a_t\)
これを \(i=0..N-1\) で足せば全ペアの差の総和になります。
具体例:\([1, 4, 6]\) のとき
- \(i=0\): 寄与 \(0\)
- \(i=1\): \(1\cdot 4 - (1)=3\)(ペア(1,4))
- \(i=2\): \(2\cdot 6 - (1+4)=7\)(ペア(1,6),(4,6)の合計)
合計 \(3+7=10\) で、\(|1-4|+|1-6|+|4-6|=3+5+2=10\) と一致します。
アルゴリズム
- 入力を次元ごとに配列
dims[k](都市 \(1..N\) の第 \(k\) 次元の値)として集める。 - 各次元 \(k=1..M\) について以下を行う:
dims[k]を昇順ソートしてarrを得る。prefを「これまで見た要素の総和(累積和)」として持つ。arrを左から順に見て、位置iの値xの寄与
\( x\cdot i - \text{pref} \) を加算する(prefは更新)。- こうして得た値が、その次元における \(\sum_{i<j}|A_{i,k}-A_{j,k}|\) になる。
- 次元ごとの結果を全て足して出力する。
計算量
- 時間計算量: \(O\big(M \cdot N \log N\big)\)(各次元でソートが必要)
- 空間計算量: \(O(MN)\)(次元ごとの値を保持)
実装のポイント
次元ごとに独立:マンハッタン距離は各次元の和なので、次元ごとに計算して足すだけで良いです。
ソート後は絶対値が外れる:昇順にしておけば \(i<j\) で \(|a_j-a_i|=a_j-a_i\) となり扱いやすくなります。
累積和
prefの使い方:s += x*i - prefをした後にpref += xと更新します。入力が大きいので高速入力:コードでは
sys.stdin.buffer.read()を使ってまとめて読み取っています。64bitに注意:Python の
intは桁あふれしませんが、他言語ならlong long等が必要です。ソースコード
import sys
def main():
it = iter(map(int, sys.stdin.buffer.read().split()))
N = next(it)
M = next(it)
dims = [[] for _ in range(M)]
for _ in range(N):
for k in range(M):
dims[k].append(next(it))
total = 0
for k in range(M):
arr = sorted(dims[k])
pref = 0
s = 0
for i, x in enumerate(arr):
s += x * i - pref
pref += x
total += s
print(total)
if __name__ == "__main__":
main()
この解説は gpt-5.2-high によって生成されました。
投稿日時:
最終更新: