概要

与えられた数列から、差の絶対値がちょうど \(K\) であるような異なる2つの要素のペアの数を求めます。

考察

この問題では、数列から2つの要素 \(A_i, A_j\) を選び、\(|A_i - A_j| = K\) となるペアの数を数える必要があります。

素朴な解法とその問題点

最も単純な方法は、すべてのペア \((i, j)\) （\(i < j\)）に対して \(|A_i - A_j| = K\) かどうかを調べることです。しかし、これは全探索であり、時間計算量が \(O(N^2)\) となり、制約 \(N \leq 2 \times 10^5\) のもとでは到底間に合いません。

また、重複要素がある場合、単純な集合や辞書による管理だけでは正確なカウントが難しくなります。

改善策：データ構造の活用

ここで重要な観察があります： - 条件 \(|A_i - A_j| = K\) は、\(A_j = A_i + K\) または \(A_j = A_i - K\) と書き換えられる。 - よって、ある要素 \(A_i\) を処理しているときに、「すでに見た要素の中に \(A_i - K\) または \(A_i + K\) が何個あるか」を高速に求められれば良い。

この「すでに見た要素のうち、ある値を持つものの個数」を高速に求めるために、Fenwick Tree（Binary Indexed Tree; BIT）を使います。

さらに、要素の値の範囲が非常に大きい（最大 \(10^9\)）ため、そのままBITのインデックスとして使うことはできません。そこで 座標圧縮 を行い、出現する値を小さい順に番号付けします。

アルゴリズム

1. 座標圧縮

   • 出現する値を一意にソートし、それぞれに \(0\) から始まるインデックスを割り振る。

2. Fenwick Tree（BIT）の準備

   • 座標圧縮後のインデックスを使って、各要素の出現回数を管理するBITを用意する。

3. 要素を順に処理

   • 各要素 \(a\) に対して以下を行う：
     • \(a - K\) および \(a + K\) が既出であれば、その個数をBITから取得して答えに加算。
     • 自身 \(a\) の出現をBITに記録（カウントを増やす）。

4. 結果を出力

このようにすることで、各要素について定数回のBIT操作で処理でき、全体で高速に解を得ることができます。

計算量

• 時間計算量: \(O(N \log N)\)
  
  • 座標圧縮に \(O(N \log N)\)、各要素ごとのBIT操作が \(O(\log N)\) で、\(N\) 回繰り返す。
• 空間計算量: \(O(N)\)
  
  • BIT、座標圧縮用のマップ、カウント用の辞書などで線形のメモリを使用。

実装のポイント

• 座標圧縮を行うことで、大きな値を小さなインデックスに対応させる。
• BITは1-indexedで扱うため、要素のインデックスに注意。
• 同じ値の要素が複数ある可能性があるため、カウントを保持する辞書も併用する。
• 参照→更新の順に行い、自分自身とのペアを誤ってカウントしないようにする。

## ソースコード

```python
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:]))

    # 座標圧縮
    sorted_unique = sorted(set(A))
    index_map = {v: i for i, v in enumerate(sorted_unique)}
    max_index = len(sorted_unique)

    # Fenwick Tree / Binary Indexed Tree
    class BIT:
        def __init__(self, size):
            self.size = size
            self.tree = [0] * (size + 1)
        
        def add(self, i, x):
            while i <= self.size:
                self.tree[i] += x
                i += i & -i
        
        def sum(self, i):
            s = 0
            while i > 0:
                s += self.tree[i]
                i -= i & -i
            return s

    bit = BIT(max_index)

    # 配列の要素をBITに追加しながら処理
    count = {}
    result = 0

    for a in A:
        # a 未満の要素の個数を取得 (a - K 未満の要素は不要なので)
        idx = index_map[a] + 1
        target_minus = a - K
        target_plus = a + K

        if target_minus in index_map:
            idx_target = index_map[target_minus] + 1
            result += bit.sum(idx_target) - bit.sum(idx_target - 1) if idx_target <= max_index else 0

        if target_plus in index_map:
            idx_target = index_map[target_plus] + 1
            result += bit.sum(idx_target) - bit.sum(idx_target - 1) if idx_target <= max_index else 0

        # 重複カウント防止のために、まず参照してから追加
        current_count = count.get(a, 0)
        count[a] = current_count + 1
        bit.add(index_map[a] + 1, 1)

    print(result)

if __name__ == "__main__":
    main()

この解説は qwen3-coder-480b によって生成されました。