概要

配列 \(A\) の各要素 \(A_i\) について、自分以外の要素との差の絶対値がちょうど \(K\) になるものが存在するかどうかを判定し、条件を満たす要素の個数を求める問題です。

考察

まず、素朴なアプローチとして「すべての生徒のペア \((i, j)\) を調べる」という方法が思いつきます。しかし、この方法では調べるペアの数が約 \(\frac{N^2}{2}\) 通りとなり、\(N = 2 \times 10^5\) の制約下では計算回数が数百億回に達してしまうため、実行時間制限（TLE）に引っかかってしまいます。

そこで、各生徒 \(i\) の視点に立って「自分と差が \(K\) になる点数」が何点かを直接計算し、その点数を取った人が全体の中にいるかを探すアプローチに切り替えます。 生徒 \(i\) の点数が \(A_i\) のとき、探すべき相手の点数は \(A_i - K\) または \(A_i + K\) のいずれかです。例えば \(A_i = 5, K = 3\) なら、点数が \(2\) または \(8\) の生徒が他にいるかを探せばよいことになります。

この「特定の点数が存在するか」という検索は、あらかじめ点数のデータを集合（Set）や連想配列（ハッシュマップ、辞書）に格納しておくことで、1回あたり \(O(1)\) と非常に高速に行うことができます。

ただし、\(K = 0\) の場合は特別な注意が必要です。 \(K = 0\) のとき、探すべき点数は \(A_i - 0 = A_i\) となり、自分自身の点数と同じになります。これを単純に集合（Set）で探してしまうと、「自分自身」を見つけて条件を満たしていると勘違いしてしまいます。そのため、\(K = 0\) の場合は「自分以外の同じ点数の生徒がいるか」、すなわち「全体の中でその点数が2回以上出現しているか」を調べる必要があります。

アルゴリズム

\(K = 0\) の場合と \(K > 0\) の場合で処理を分けます。

1. \(K = 0\) の場合 1. 連想配列（Pythonの dict など）を用意し、各点数が何回出現したかをカウントします。 2. すべての生徒 \(i\) について、自分の点数 \(A_i\) の出現回数が \(2\) 回以上であれば、自分以外にも同じ点数の人がいるということなので、答えのカウントを \(1\) 増やします。

2. \(K > 0\) の場合 1. 配列 \(A\) のすべての要素を集合（Pythonの set など）に格納します。 2. すべての生徒 \(i\) について、集合の中に \(A_i - K\) または \(A_i + K\) が存在するかを確認します。存在すれば、答えのカウントを \(1\) 増やします。 （\(K > 0\) であれば \(A_i \pm K\) が \(A_i\) 自身になることはないため、自分自身を誤って見つけてしまう心配はありません）

計算量

• 時間計算量: \(O(N)\)
  • 連想配列や集合へのデータの追加に \(O(N)\)、各生徒に対する検索が \(O(1)\) でそれを \(N\) 回繰り返すため \(O(N)\) となります。全体として \(O(N)\) であり、制限時間内に余裕で処理できます。
• 空間計算量: \(O(N)\)
  • 登場する点数の種類を連想配列や集合に保存するため、最大で \(O(N)\) のメモリを使用します。

実装のポイント

• データの検索速度: Pythonのリスト（配列）に対する in 演算子を使った検索は \(O(N)\) かかってしまうため、必ず set や dict を使って検索を \(O(1)\) に高速化することが重要です。

• 条件を満たす「生徒」の数を数える: 問題文の「条件を満たすペアの組数ではなく、条件を満たす生徒が全部で何人いるかを答える」という点に注意します。上記のアルゴリズムでは、各生徒 \(i\) を主語にして判定を行っているため、自然と生徒の人数を数えることができます。

  ソースコード

import sys

def main():
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    
    N = int(input_data[0])
    K = int(input_data[1])
    A = [int(x) for x in input_data[2:N+2]]
    
    ans = 0
    if K == 0:
        counts = {}
        for a in A:
            counts[a] = counts.get(a, 0) + 1
        for a in A:
            if counts[a] >= 2:
                ans += 1
    else:
        S = set(A)
        for a in A:
            if (a - K) in S or (a + K) in S:
                ans += 1
                
    print(ans)

if __name__ == '__main__':
    main()

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