Official

B - ペアを見つけよう / Find a Pair Editorial by admin Editorial for printing

Qwen3-Coder-480B

概要

各生徒の点数から、自分以外に点数の差がちょうど \(K\) となる生徒がいる人数を数える問題です。

考察

この問題では、各生徒 \(i\) に対して、以下のような単純な方法で調べると効率が悪くなります:

  • 生徒 \(i\) に対して、他の全生徒 \(j\)(\(j \neq i\))と点数の差を一つずつ比較する → 計算量は \(O(N^2)\) となり、\(N\) が最大 \(2 \times 10^5\) なので間に合わない(TLE)。

そこで、以下のような工夫が必要になります:

  • 各点数が何回出現したかを記録しておけば、生徒の個別比較を避けることができる。
  • 特に \(K = 0\) の場合は、同じ点数の生徒が他にいるかどうかだけを気にすれば良い。
  • \(K > 0\) の場合は、ある点数 \(x\) に対して \(x - K\) または \(x + K\) の点数を持つ生徒が存在するかどうかを確認すれば良い。

このように、点数の種類ごとにまとめることで、無駄な比較を省くことができます。

アルゴリズム

  1. 最初に、各点数が何回出現したかを collections.Counter を使ってカウントします。
  2. \(K = 0\) の場合:
    • 同じ点数を持つ生徒が2人以上いれば、その点数を持つ全員が条件を満たすので、その人数を加算します。
  3. \(K > 0\) の場合:
    • 各点数 \(x\) について、\(x - K\) または \(x + K\) が存在するかどうかをセットで高速に調べる。
    • 存在すれば、点数 \(x\) を持つ全員が条件を満たすので、その人数を加算します。

例

入力:

5 2
1 3 3 5 7

点数のカウント: {1: 1, 3: 2, 5: 1, 7: 1}
点数の集合: {1, 3, 5, 7}

  • 点数 1: \(1+2=3\) が存在 → 条件を満たす(1人)
  • 点数 3: \(3-2=1\) と \(3+2=5\) が両方存在 → 条件を満たす(2人)
  • 点数 5: \(5-2=3\) が存在 → 条件を満たす(1人)
  • 点数 7: \(7-2=5\) が存在 → 条件を満たす(1人)

合計: \(1 + 2 + 1 + 1 = 5\) 人

計算量

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

点数の種類数は最大でも \(N\) 個なので、カウントや集合操作がすべて \(O(N)\) で行えます。

実装のポイント

  • Counter を使うことで、点数ごとの出現回数を簡単に取得できる。

  • set を使うことで、特定の点数が存在するかを \(O(1)\) で判定できる。

  • \(K = 0\) の場合とそうでない場合で処理を分ける必要がある。

    ソースコード

from collections import Counter

N, K = map(int, input().split())
A = list(map(int, input().split()))

counter = Counter(A)
ans = 0

if K == 0:
    # 同じ点数を持つ人が2人以上いる場合、その点数を持つ全員が条件を満たす
    for score, count in counter.items():
        if count >= 2:
            ans += count
else:
    # K > 0 の場合、Ai - K または Ai + K が存在するかどうかをチェック
    unique_scores = set(A)
    for score in unique_scores:
        if (score - K) in unique_scores or (score + K) in unique_scores:
            ans += counter[score]

print(ans)

この解説は qwen3-coder-480b によって生成されました。

posted:
last update: