概要

平面上の \(N\) 個の集落から、指定された集落 \(C_k\) への「切り捨てユークリッド距離」の総和を求めるクエリに \(Q\) 回答える問題です。クエリごとに計算すると間に合わないため、すべての集落のペア間の距離を事前に計算しておく（前計算）ことで、各クエリに \(O(1)\) で高速に回答します。

考察

素朴なアプローチとその限界

最も愚直な方法は、各クエリ \(C_k\) が与えられるたびに、他のすべての集落 \(j\) との距離 \(\lfloor \sqrt{(X_{C_k} - X_j)^2 + (Y_{C_k} - Y_j)^2} \rfloor\) を計算して足し合わせる方法です。 この場合、1つのクエリあたり \(O(N)\) の計算量がかかります。クエリ数は \(Q\) 個あるため、全体の時間計算量は \(O(NQ)\) となります。

制約を確認すると、 \(N \le 2000\)、 \(Q \le 2 \times 10^5\) であるため、最悪の場合の計算回数は： $\(N \times Q = 2000 \times (2 \times 10^5) = 4 \times 10^8\)$ となり、一般的な制限時間（2秒）の中では実行時間制限超過（TLE）になってしまいます。

高速化の鍵：クエリの独立性と前計算

ここで、クエリで聞かれるのは「集落 \(C_k\) に基地を置いたときの距離の総和」です。この値はクエリの順番や他のクエリの内容に依存せず、集落の配置だけで一意に決まります。また、集落の数 \(N\) は最大でも \(2000\) と比較的小さいです。

したがって、あらかじめすべての集落 \(i\) \((1 \leq i \leq N)\) について、基地を \(i\) に置いたときの距離の総和を計算（前計算）しておくことを考えます。

集落の全ペア \((i, j)\) の組み合わせは \(\frac{N(N-1)}{2}\) 通りです。 \(N = 2000\) のとき、このペア数は： $\(\frac{2000 \times 1999}{2} \approx 2 \times 10^6\)$ となり、コンピュータにとっては一瞬で計算できる規模です。

集落 \(i\) と集落 \(j\) の間の距離は、基地を \(i\) に置いたときも \(j\) に置いたときも同じ（対称性：\(dist(i, j) = dist(j, i)\)）なので、1回の距離計算の結果を両方の集落の合計値に加算することができます。これにより、無駄な計算を半分に減らすことができます。

アルゴリズム

1. 配列の準備: 各集落 \(i\) についての距離の総和を記録する配列 ans を長さ \(N\) で用意し、 \(0\) で初期化します。

2. すべてのペアに対する前計算: 二重ループを用いて、すべての集落のペア \((i, j)\) （ただし \(i < j\)）について以下を行います。

   • 集落 \(i\) と集落 \(j\) のユークリッド距離 \(d = \lfloor \sqrt{(X_i - X_j)^2 + (Y_i - Y_j)^2} \rfloor\) を計算する。
   • ans[i] と ans[j] のそれぞれに \(d\) を加算する。

3. クエリへの回答: 各クエリ \(C_k\) に対して、前計算した配列から ans[C_k - 1] を取り出して出力します（0-indexedに調整します）。

計算量

• 時間計算量: \(O(N^2 + Q)\)

  • 前計算において、すべてのペアの距離を計算するのに \(O(N^2)\) 時間かかります。
  • 各クエリには \(O(1)\) で答えられるため、 \(Q\) 回のクエリ処理全体で \(O(Q)\) 時間かかります。
  • \(N \le 2000, Q \le 2 \times 10^5\) のとき、全体の計算回数は \(2 \times 10^6 + 2 \times 10^5 \approx 2.2 \times 10^6\) 回となり、制限時間内に余裕で間に合います。

• 空間計算量: \(O(N + Q)\)

  • 各集落の座標を保持する配列に \(O(N)\)、前計算の結果を保持する配列に \(O(N)\)、クエリの出力を保持する配列に \(O(Q)\) のメモリを使用します。

実装のポイント

• 高速な入出力: Pythonにおいて \(Q = 2 \times 10^5\) のような大量の入出力を扱う場合、input() や print() を愚直に繰り返すと入出力自体がボトルネックになります。sys.stdin.read で入力を一括で読み込み、sys.stdout.write で一括で出力することで、実行時間を大幅に短縮できます。

• 平方根の切り捨て: \(\lfloor \sqrt{x} \rfloor\) は、Pythonでは int(math.sqrt(x)) を用いて簡潔かつ高速に計算できます。

  ソースコード

import math
import sys


def solve():
    # すべての入力を高速に読み込む
    input_data = sys.stdin.read().split()
    if not input_data:
        return

    N = int(input_data[0])
    Q = int(input_data[1])

    X = [0] * N
    Y = [0] * N
    idx = 2
    for i in range(N):
        X[i] = int(input_data[idx])
        Y[i] = int(input_data[idx + 1])
        idx += 2

    # 各集落から他のすべての集落への切り捨て距離の合計を前計算する
    # 対称性 (dist(i, j) == dist(j, i)) を利用してループ回数を半分にする
    ans = [0] * N
    sqrt = math.sqrt

    for i in range(N):
        xi = X[i]
        yi = Y[i]
        for j in range(i + 1, N):
            dx = xi - X[j]
            dy = yi - Y[j]
            d = int(sqrt(dx * dx + dy * dy))
            ans[i] += d
            ans[j] += d

    # 各クエリに対して、前計算した結果を O(1) で取得して出力する
    out = []
    for _ in range(Q):
        c = int(input_data[idx]) - 1
        idx += 1
        out.append(str(ans[c]))

    sys.stdout.write("\n".join(out) + "\n")


if __name__ == "__main__":
    solve()

この解説は gemini-3.5-flash-thinking によって生成されました。