C++で配列内の条件を満たすペアの数を効率的に求めるプログラム

問題の概要

配列numsにn個の整数が与えられていると仮定します。この中から2つの数値を選んでペアを作りますが、その際「2つの要素の配列内における位置の差が、2つの数値の合計と等しい」という条件を満たす必要があります。n個の要素から選べるペアの総数は n(n − 1) / 2 通りであり、その中から条件を満たすペアの個数を求めるのが本記事のテーマです。

例えば、入力が n = 8、nums = {4, 2, 1, 0, 1, 2, 3, 3} である場合、出力は 13 になります。つまり、この配列には条件を満たすペアが13個存在します。

解法のアプローチ

すべてのペアを総当たりで確認する方法もありますが、以下の手順に従うことで、二分探索を利用した効率的な解法を実現できます。

長さnの配列valsを用意する
i := 0 から i < n の間、i を1ずつ増やしながら繰り返す：
    vals[i] := i + 1 - nums[i]
配列valsを昇順にソートする
res := 0 と初期化する
i := 0 から i < n の間、i を1ずつ増やしながら繰り返す：
    k := nums[i] + i + 1
    res := res + （vals内でk以上の値が最初に現れる位置 − vals内でkより大きい値が最初に現れる位置）
res を返す

ここで重要なのは、「位置の差 ＝ 2つの数の和」という条件を変形すると、ペア (i, j) に対して j − nums[j] ＝ i + nums[i] という関係が成り立つ点です。各要素についてあらかじめ i + 1 − nums[i] を計算して配列valsに格納し、ソートしておけば、upper_bound と lower_bound の差（＝vals内でkと等しい要素の個数）を二分探索により O(log n) で求められます。

C++による実装例

理解を深めるために、実際の実装コードを見てみましょう。

#include <bits/stdc++.h>
using namespace std;

int solve(int n, vector<int> nums){
    vector<int> vals(n);
    for( int i = 0; i < n; i++)
        vals[i] = i + 1 - nums[i];
    sort(vals.begin(), vals.end());
    int res = 0;
    for( int i = 0; i < n; i++ ) {
        int k = nums[i] + i + 1;
        res += upper_bound(vals.begin(), vals.end(), k) - lower_bound(vals.begin(), vals.end(), k);
    }
    return res;
}
int main() {
    int n = 8;
    vector<int> nums = {4, 2, 1, 0, 1, 2, 3, 3};
    cout<< solve(n, nums);
    return 0;
}

入力

8, {4, 2, 1, 0, 1, 2, 3, 3}

出力

13

計算量の評価

ソートに O(n log n)、さらに各要素ごとの二分探索にも O(log n) かかるため、全体の時間計算量は O(n log n) となります。すべてのペアを順に調べる総当たり法（O(n²)）と比較すると、大規模な入力に対して格段に高速に動作する点が大きなメリットです。