概要

\(N\) 個の整数からなる数列 \(A\) の中から、差の絶対値がちょうど \(K\) となる 2 つの要素の組み合わせ（ペア）が何通りあるかを求める問題です。

考察

素朴なアプローチ

最も単純な方法は、すべての 2 要素の組み合わせ \((A_i, A_j)\) を二重ループで確認することです。しかし、この方法では組み合わせの数が \(\frac{N(N-1)}{2}\) 通りとなり、計算量は \(O(N^2)\) です。 制約では \(N \le 2 \times 10^5\) であるため、\(N^2\) は約 \(4 \times 10^{10}\) に達し、実行制限時間に間に合いません（TLEとなります）。

効率的なアプローチ

差の絶対値が \(K\) であるということは、ある要素 \(x\) に対して、ペアの相手は \(x + K\) または \(x - K\) のいずれかであることを意味します。 ここで、すべての \(x\) について「\(x\) と \(x + K\)」というペアのみを数え上げることにすれば、重複なくすべてのペアを網羅できます（\(K \ge 1\) なので \(x\) と \(x+K\) は必ず異なります）。

具体的には、以下の手順で考えます。 1. 数列に含まれる各数値が「何回出現するか」をあらかじめ記録しておく。 2. 数列に含まれる各数値 \(x\) について、その出現回数を \(C_x\) とする。 3. \(x + K\) という数値が数列に \(C_{x+K}\) 回出現する場合、それらの組み合わせは \(C_x \times C_{x+K}\) 通り存在する。 4. すべての \(x\) についてこの値を合計する。

この方法なら、各数値を 1 回ずつ確認するだけで済むため、高速に処理が可能です。

アルゴリズム

ハッシュマップ（Pythonでは dict や collections.Counter）を利用した解法を用います。

1. カウント: 数列 \(A\) を走査し、各数値の出現回数をハッシュマップに格納します。
2. 集計: ハッシュマップに含まれる各キー \(x\) について以下を繰り返します。
   • もしハッシュマップに \(x + K\) というキーが存在すれば、(xの出現回数) × (x + Kの出現回数) を答えに加算する。
3. 出力: 最終的な合計値を出力する。

計算量

• 時間計算量: \(O(N)\)
  • 数列の読み込みに \(O(N)\)、ハッシュマップの作成に \(O(N)\)、ハッシュマップの各要素の確認に最大 \(O(N)\) かかります。ハッシュマップの操作（挿入・検索）は平均 \(O(1)\) であるため、全体で \(O(N)\) となります。
• 空間計算量: \(O(N)\)
  • 各数値を格納するハッシュマップのサイズが最大で \(N\) となるため、\(O(N)\) のメモリを消費します。

実装のポイント

• Counterクラス: Pythonの collections.Counter を使うと、リストから要素の出現回数を数え上げる処理を簡潔に書くことができます。

• 重複カウントの回避: ループ内で x + K だけをチェックすることで、「小さい方の値から大きい方の値を探す」という一方向の確認になり、同じペアを 2 回数えてしまうのを防いでいます。

• 高速な入力: \(N\) が大きいため、sys.stdin.read().split() などを用いて一括で入力を読み込むと、実行時間を短縮できます。

  ソースコード

import sys
from collections import Counter

def solve():
    # 入力を標準入力から一括で読み込む（高速化のため）
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    
    # N: 数列の要素数, K: 目標とする差
    N = int(input_data[0])
    K = int(input_data[1])
    # A: 数列の各要素
    A = list(map(int, input_data[2:]))
    
    # 各数値の出現回数をカウントする
    counts = Counter(A)
    
    ans = 0
    # 数列に含まれるユニークな各数値 x についてループ
    for x in counts:
        # x + K が数列内に存在するか確認する
        # x + K を探すことで、すべてのペア {x, x + K} を重複なく数え上げることができる
        # (K >= 1 なので、x と x + K は必ず異なる値になる)
        target = x + K
        if target in counts:
            # x の出現回数と target の出現回数の積を足す
            ans += counts[x] * counts[target]
            
    # 結果を出力
    print(ans)

if __name__ == "__main__":
    solve()

この解説は gemini-3-flash-preview によって生成されました。