概要

\(N\) 人の生徒の中で、自分以外に点数の差の絶対値がちょうど \(K\) である相手が少なくとも1人いる生徒の人数を求める問題です。

考察

素朴なアプローチとその問題点

最も単純な方法は、すべての生徒のペア \((i, j)\) を調べて \(|A_i - A_j| = K\) かどうかを確認することです。しかし、ペアの数は \(\frac{N(N-1)}{2}\) 通りあり、\(N\) が最大 \(2 \times 10^5\) のとき約 \(2 \times 10^{10}\) 回の計算が必要になるため、TLE（時間超過） になります。

重要な気づき

\(|A_i - A_j| = K\) という条件は、次のように言い換えられます：

\[A_j = A_i + K \quad \text{または} \quad A_j = A_i - K\]

つまり、生徒 \(i\) の点数が \(A_i\) のとき、点数が \(A_i + K\) または \(A_i - K\) の生徒が他に存在するかを調べればよいのです。

これは各点数の出現回数をハッシュマップ（辞書）で管理すれば、1人あたり \(O(1)\) で判定できます。

\(K = 0\) の場合の注意

\(K = 0\) の場合、\(A_j = A_i\) となる \(j \neq i\) が存在するかを調べます。つまり、同じ点数の生徒が2人以上いるかどうかがポイントです。同じ点数が1人しかいなければ条件を満たしません。

アルゴリズム

1. 各点数の出現回数を Counter（ハッシュマップ）で数える。
2. \(K = 0\) の場合: 出現回数が \(2\) 以上の点数について、その出現回数をすべて合計する。
3. \(K \neq 0\) の場合: 各点数 \(v\) について、\(v + K\) または \(v - K\) が Counter に存在するかを調べる。存在すれば、点数 \(v\) の生徒は全員条件を満たすので、その出現回数を合計する。

具体例

\(N = 5\), \(K = 2\), \(A = [3, 1, 5, 3, 7]\) の場合：

• 出現回数: \(\{3: 2, 1: 1, 5: 1, 7: 1\}\)
• 点数 \(3\): \(3+2=5\) が存在 → 条件を満たす（2人）
• 点数 \(1\): \(1+2=3\) が存在 → 条件を満たす（1人）
• 点数 \(5\): \(5-2=3\) が存在 → 条件を満たす（1人）
• 点数 \(7\): \(7-2=5\) が存在 → 条件を満たす（1人）
• 答え: \(2 + 1 + 1 + 1 = 5\) 人

計算量

• 時間計算量: \(O(N)\)（Counter の構築に \(O(N)\)、各点数の判定に \(O(1)\) × 点数の種類数）
• 空間計算量: \(O(N)\)（Counter の格納に必要）

実装のポイント

• \(K = 0\) の場合を分けて処理する: \(K = 0\) のとき \(v + K = v - K = v\) なので、「自分自身と同じ点数」を探すことになりますが、自分以外に同じ点数の人がいる必要があるため、出現回数が \(2\) 以上かどうかで判定します。

• \(K \neq 0\) の場合は \(v + K\) と \(v - K\) の両方を調べる: 絶対値の条件なので、プラス方向とマイナス方向の両方をチェックする必要があります。

• ペアの組数ではなく人数を答える: 条件を満たすペアが複数あっても、その生徒は1人として数えます。Counter を使って点数の種類ごとにまとめて処理することで、自然にこの要件を満たせます。

  ソースコード

import sys
from collections import Counter

def main():
    input_data = sys.stdin.read().split()
    N = int(input_data[0])
    K = int(input_data[1])
    A = list(map(int, input_data[2:2+N]))
    
    cnt = Counter(A)
    
    ans = 0
    if K == 0:
        for v, c in cnt.items():
            if c >= 2:
                ans += c
    else:
        for v, c in cnt.items():
            if (v + K) in cnt or (v - K) in cnt:
                ans += c
    
    print(ans)

main()

---

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