概要

\(N\) 個の集落から基地となる集落を選び、その基地から全集落へのユークリッド距離（各々切り捨て）の合計を \(Q\) 回の検討それぞれについて求める問題です。

考察

素朴なアプローチ

各クエリについて、基地から全 \(N\) 個の集落への距離を計算すると、1回のクエリあたり \(O(N)\) の計算が必要です。\(Q\) 回のクエリすべてに対して愚直に計算すると \(O(NQ)\) となります。

\(N \leq 2000\), \(Q \leq 2 \times 10^5\) なので \(NQ\) は最大 \(4 \times 10^8\) となり、そのまま実装するとTLEの危険があります。

重要な気づき：クエリの重複排除

基地は必ずいずれかの集落（\(1\) ～ \(N\)）に設置されます。したがって、異なる基地の候補は高々 \(N\) 個しかありません。同じ集落が複数回クエリで指定されても、答えは同じです。

そこで、ユニークなクエリ（基地の候補）についてのみ計算を行い、結果をキャッシュすれば、計算量は \(O(N^2)\)（最大で \(N\) 種類の基地 × 各 \(N\) 集落への距離計算）に抑えられます。

\(N \leq 2000\) なので \(N^2 = 4 \times 10^6\) 程度となり、十分高速です。

アルゴリズム

1. 全集落の座標を読み込む。
2. 全クエリを読み込み、ユニークな基地候補の集合を求める。
3. 各ユニークな基地候補 \(c\) について、全集落 \(j\) への距離の合計を計算する： $\(\text{results}[c] = \sum_{j=1}^{N} \lfloor \sqrt{(X_c - X_j)^2 + (Y_c - Y_j)^2} \rfloor\)\( ここで整数平方根（`isqrt`）を使うことで、浮動小数点誤差なく \)\lfloor \sqrt{d^2} \rfloor$ を正確に計算できます。
4. 各クエリに対して、事前計算した結果を出力する。

具体例

例えば集落が \((0, 0)\), \((3, 4)\), \((1, 1)\) の3つで、基地を集落1（座標 \((0,0)\)）に置く場合： - 集落1への距離: \(\lfloor\sqrt{0}\rfloor = 0\) - 集落2への距離: \(\lfloor\sqrt{9+16}\rfloor = \lfloor 5 \rfloor = 5\) - 集落3への距離: \(\lfloor\sqrt{1+1}\rfloor = \lfloor\sqrt{2}\rfloor = 1\)

合計は \(0 + 5 + 1 = 6\) となります。

計算量

• 時間計算量: \(O(N^2 + Q)\)
  • ユニークな基地候補は最大 \(N\) 個、各候補について \(N\) 集落への距離を計算するので \(O(N^2)\)
  • クエリの出力に \(O(Q)\)
• 空間計算量: \(O(N + Q)\)
  • 集落の座標に \(O(N)\)、クエリの格納と結果のキャッシュに \(O(N + Q)\)

実装のポイント

• isqrt の使用: 浮動小数点の math.sqrt を使うと丸め誤差で正しい床関数の値が得られない場合があります。Python の math.isqrt は整数の平方根を正確に切り捨てで返すため、\(\lfloor\sqrt{d^2}\rfloor\) を誤差なく計算できます。

• クエリの重複排除: set を使ってユニークなクエリのみ計算することで、\(Q\) が大きくても計算量を \(O(N^2)\) に抑えられます。

• 高速な入出力: sys.stdin.buffer.read() で一括読み込みし、出力も '\n'.join() でまとめて書き出すことで、Python でもTLEを回避しやすくなります。

  ソースコード

import sys
import math

def main():
    input_data = sys.stdin.buffer.read().split()
    idx = 0
    N = int(input_data[idx]); idx += 1
    Q = int(input_data[idx]); idx += 1
    
    X = [0] * N
    Y = [0] * N
    for i in range(N):
        X[i] = int(input_data[idx]); idx += 1
        Y[i] = int(input_data[idx]); idx += 1
    
    # Precompute answers for each unique base
    # N, Q <= 2000, so O(N*N) = 4*10^6 which is fine
    # But queries may repeat, so we cache by village index
    
    from math import isqrt
    
    results = {}
    
    queries = []
    for k in range(Q):
        c = int(input_data[idx]) - 1; idx += 1
        queries.append(c)
    
    unique_queries = set(queries)
    
    for c in unique_queries:
        cx, cy = X[c], Y[c]
        total = 0
        for j in range(N):
            dx = cx - X[j]
            dy = cy - Y[j]
            d2 = dx * dx + dy * dy
            total += isqrt(d2)
        results[c] = total
    
    out = []
    for c in queries:
        out.append(str(results[c]))
    
    sys.stdout.write('\n'.join(out) + '\n')

main()

---

この解説は claude4.6opus-thinking によって生成されました。