概要

数列から異なる 2 要素を選び、差の絶対値がちょうど \(K\) になるペア数を、重複要素も考慮して高速に数えます。

考察

素朴に全てのペア \((i, j)\)（\(i<j\)）を調べると、組み合わせ数は \(\frac{N(N-1)}{2}\) なので最大で約 \(2\times 10^5\) に対して \(O(N^2)\) となり、確実に時間切れになります。

ここで重要な観察は次の通りです：

• 「差の絶対値が \(K\)」とは、値の組としては \((x, x+K)\) の形を数えれば十分です。
  なぜなら、\(|a-b|=K\) となるのは \(a=b+K\) または \(b=a+K\) のどちらかで、結局「小さい方を \(x\)、大きい方を \(x+K\)」と一意に定まるからです。
• 同じ値が複数回出てくるので、個数（頻度）を使うとまとめて数えられます。
  例えば \(x\) が \(c\) 回、\(x+K\) が \(d\) 回出るなら、差が \(K\) になるペアは \(c \times d\) 個あります（それぞれの出現位置の組み合わせ）。

具体例：\(A=[1,1,4,4,4]\), \(K=3\)
- \(1\) は 2 回、\(4\) は 3 回なので、\((1,4)\) のペアは \(2\times 3=6\) 個。

この方針なら、各値 \(x\) について「\(x+K\) が何個あるか」を高速に取得できればよく、ハッシュ（辞書）で頻度管理すれば \(O(N)\) 近くで解けます。

アルゴリズム

1. 数列 \(A\) の各値の出現回数を Counter（辞書）で数える。cnt[v] = v の個数。
2. 答え ans=0 を用意する。
3. cnt に含まれる各値 \(x\) について、
   • \(x\) の個数を \(c\) とし、\(x+K\) の個数を \(d = \text{cnt}[x+K]\)（なければ 0）とする。
   • ans += c * d を加算する。
4. ans を出力する。

この数え方では常に「小さい方 \(x\) と大きい方 \(x+K\)」だけを数えるため、同じペアを 2 回数えることがありません（\((x, x+K)\) と \((x+K, x)\) を別扱いしないという条件を自然に満たします）。

計算量

• 時間計算量: \(O(N)\)（Counter 構築が \(O(N)\)、異なる値の種類数を \(M\) とするとループは \(O(M)\) で \(M \le N\)）
• 空間計算量: \(O(M)\)（値の種類数ぶんの頻度を保持）

実装のポイント

• 頻度管理には collections.Counter が便利です。

• cnt.get(x + K, 0) のように書くことで、存在しないキーでも 0 を返せて安全です。

• ans は最大で \(O(N^2)\) 規模になり得ますが、Python の整数はオーバーフローしないのでそのままで問題ありません。

  ソースコード

import sys
from collections import Counter

def main():
    input = sys.stdin.readline
    N, K = map(int, input().split())
    A = list(map(int, input().split()))
    cnt = Counter(A)

    ans = 0
    for x, c in cnt.items():
        ans += c * cnt.get(x + K, 0)

    print(ans)

if __name__ == "__main__":
    main()

この解説は gpt-5.2-high によって生成されました。