概要

\(N\) 個の集落の中から基地となる集落を 1 つ選び、そこから全集落への「ユークリッド距離の切り捨て値」の合計を求める問題です。検討（クエリ）が \(Q\) 回行われるため、効率的な計算が求められます。

考察

1. 素朴な解法とその限界

各クエリに対して、毎回すべての集落（\(N\) 個）との距離を計算すると、全体の計算量は \(O(Q \times N)\) となります。 本問題の制約は \(N \le 2000, Q \le 2 \times 10^5\) であるため、計算回数は最大で \(2 \times 10^5 \times 2000 = 4 \times 10^8\) 回に達します。一般的な制限時間（2秒程度）において、特に Python ではこの計算量を処理しきれず、実行時間制限（TLE）になってしまいます。

2. 重要な気づき：基地の候補は限られている

基地として選ばれる場所は、必ず \(N\) 個の集落のいずれかです。つまり、「どの集落を基地にするか」というパターンは最大でも \(N\) 通りしかありません。 同じ集落が何度もクエリで指定される可能性があるため、一度計算した結果を再利用することで、無駄な計算を省くことができます。

3. 前計算による高速化

クエリを処理する前に、各集落を基地とした場合の合計距離をあらかじめ計算（前計算）しておきます。 - 基地の候補となる集落の座標ごとに計算を行う。 - 計算した結果を辞書（連想配列）などに保存しておく。 - 各クエリに対しては、保存しておいた結果を \(O(1)\) で取り出す。

これにより、全体の計算量を \(O(N^2 + Q)\) に抑えることができます。\(2000^2 + 2 \times 10^5 \approx 4.2 \times 10^6\) となり、制限時間内に十分間に合う計算量になります。

アルゴリズム

1. データの集約: 同じ座標に複数の集落がある場合、それらをまとめてカウントすることで計算をさらに効率化できます。座標 \((X, Y)\) ごとの集落数を Counter 等で集計します。
2. ユニークな基地座標の抽出: 全クエリの中で、基地として指定された集落の座標を重複なく抽出します。
3. 距離合計の計算（メイン処理）: 抽出した各基地座標 \((X_{base}, Y_{base})\) について、以下の手順で合計距離を求めます。
   • 全ての集落座標 \((X_j, Y_j)\) に対して、ユークリッド距離 \(d = \sqrt{(X_{base} - X_j)^2 + (Y_{base} - Y_j)^2}\) を計算する。
   • \(\lfloor d \rfloor\) を求め、その座標にある集落の数を掛けて合計に加算する。
   • 求めた合計値を、基地の座標をキーとして辞書に保存する。
4. クエリへの回答: 各クエリで与えられた集落番号 \(C_k\) の座標を特定し、辞書から計算済みの値を参照して出力します。

計算量

• 時間計算量: \(O(N^2 + Q)\)
  • 座標の集計に \(O(N)\)。
  • 前計算において、最大 \(N\) 個の基地候補に対し \(N\) 個の集落との距離を計算するため \(O(N^2)\)。
  • クエリの回答に \(O(Q)\)。
  • 実際には座標の重複がある場合、前計算は \(O((\text{異なる座標の数})^2)\) となり、より高速になります。
• 空間計算量: \(O(N + Q)\)
  • 座標データ、集計データ、およびクエリの結果を保持するために必要です。

実装のポイント

• 整数平方根の利用: \(\lfloor \sqrt{x} \rfloor\) の計算には、Python の math.isqrt(x) を使用すると、浮動小数点の精度問題を回避しつつ高速に整数部分を取得できます。

• 高速な入出力: \(Q\) が大きいため、sys.stdin.read().split() による一括読み込みや、sys.stdout.write() による一括出力を行うことで、入出力のボトルネックを解消しています。

• 座標による管理: 集落番号ではなく「座標」をキーにして計算結果を管理することで、異なる集落番号でも座標が同じであれば計算をスキップできます。

  ソースコード

import sys
import math
from collections import Counter

def solve():
    """
    救急ヘリコプターの配置問題を解く関数。
    
    N個の集落の座標を読み取り、Q個のクエリ（基地を設置する集落番号）に対して、
    基地から全集落へのユークリッド距離の切り捨ての合計を求める。
    計算量は O(D_B * D + Q) であり、ここで D は集落の異なる座標の数、
    D_B はクエリで指定された集落の異なる座標の数。
    最悪の場合でも O(N^2 + Q) となり、制約下で十分高速に動作する。
    """
    
    # 標準入力から全データを一括で読み込み、スペース区切りで分割する
    try:
        input_data = sys.stdin.read().split()
    except EOFError:
        return
    
    if not input_data:
        return
    
    # イテレータを使用してデータを順番に取得する
    it = iter(input_data)
    try:
        N = int(next(it))
        Q = int(next(it))
    except StopIteration:
        return
    
    # 全集落の座標をリストに格納する
    all_coords = []
    for _ in range(N):
        try:
            x = int(next(it))
            y = int(next(it))
            all_coords.append((x, y))
        except StopIteration:
            break
    
    # 各検討（クエリ）における基地の設置先集落番号を読み込む（0-indexedに変換）
    query_indices = []
    for _ in range(Q):
        try:
            query_indices.append(int(next(it)) - 1)
        except StopIteration:
            break
            
    # 集落の座標ごとの頻度をカウントする（同じ座標に複数の集落がある場合をまとめる）
    counts = Counter(all_coords)
    # 異なる座標、およびその座標にある集落の数をリスト化する
    distinct_info = [(x, y, count) for (x, y), count in counts.items()]
    
    # クエリで指定された集落のうち、ユニークな座標のセットを作成する
    # これにより、同じ座標にある集落を基地とする検討を一度に済ませる
    unique_query_coords = {all_coords[idx] for idx in query_indices}
    
    # 各ユニークな座標を基地とした時の合計距離を格納する辞書
    base_results = {}
    _isqrt = math.isqrt  # 整数平方根（切り捨て）関数
    
    # 各ユニークな基地座標に対して、全集落への距離の合計を計算する
    for xi, yi in unique_query_coords:
        total = 0
        for xj, yj, count_j in distinct_info:
            # ユークリッド距離の二乗を計算
            dx = xi - xj
            dy = yi - yj
            dist_sq = dx * dx + dy * dy
            # 切り捨てられたユークリッド距離を求め、その座標にある集落の数を掛ける
            total += _isqrt(dist_sq) * count_j
        base_results[(xi, yi)] = total
        
    # 各クエリの設置先集落の座標に基づき、計算済みの結果を紐付ける
    # 文字列のリストとして保持し、最後に一括で出力する
    output = [str(base_results[all_coords[idx]]) for idx in query_indices]
    
    # 全ての結果を改行で繋いで標準出力に書き出す
    sys.stdout.write('\n'.join(output) + '\n')

if __name__ == '__main__':
    solve()

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