概要

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

考察

各生徒 \(i\) について、自分以外の生徒の中に \(|A_i - A_j| = K\) となる生徒 \(j\) が存在するかどうかを調べる必要があります。素朴にすべてのペアを調べると \(O(N^2)\) の計算量となり、\(N\) が最大 \(2 \times 10^5\) なので時間制限に間に合いません。

重要な観察として、生徒 \(i\) が条件を満たすためには、\(A_i + K\) または \(A_i - K\) という点数を持つ別の生徒が少なくとも1人いればよいという点です。ただし、\(K = 0\) の場合は特別で、自分と同じ点数を持つ別の生徒が少なくとも1人いればよいことになります。

アルゴリズム

1. まず、各点数の出現回数をハッシュマップ（辞書）で数えます。
2. 各生徒 \(i\) について：
   • \(K = 0\) の場合：同じ点数の生徒が自分以外にいるか（出現回数が2以上か）をチェック
   • \(K > 0\) の場合：\(A_i + K\) または \(A_i - K\) という点数を持つ生徒が存在するかをチェック
3. 条件を満たす生徒の人数をカウントします。

計算量

• 時間計算量: \(O(N)\)
  • 点数リストの走査とハッシュマップの操作はそれぞれ \(O(N)\) です
• 空間計算量: \(O(N)\)
  • ハッシュマップに \(N\) 個の要素を格納するため

実装のポイント

• ハッシュマップ（defaultdict）を使って各点数の出現回数を効率的に管理します

• \(K = 0\) の場合と \(K > 0\) の場合で処理を分ける必要があります

• \(K\) が非常に大きい場合（\(10^9\) 程度）でも、ハッシュマップの検索は \(O(1)\) で行えるため問題ありません

• \(A_i \pm K\) が負の値になる可能性がありますが、点数は0以上の整数なので \(A_i - K\) が負の場合は自動的に存在しないと判断できます（ハッシュマップに存在しないキーを検索してもエラーになりません）

  ソースコード

def main():
    import sys
    input = sys.stdin.read
    data = input().split()
    
    n = int(data[0])
    K = int(data[1])
    A = list(map(int, data[2:2+n]))
    
    if n == 0:
        print(0)
        return
        
    from collections import defaultdict
    count_map = defaultdict(int)
    for a in A:
        count_map[a] += 1
        
    ans = 0
    for a in A:
        if K == 0:
            if count_map[a] >= 2:
                ans += 1
        else:
            if count_map[a + K] > 0 or count_map[a - K] > 0:
                ans += 1
                
    print(ans)

if __name__ == "__main__":
    main()

この解説は deepseekv3 によって生成されました。