C++で合計が最小となるK個のペアを効率的に検索する方法

問題概要

2つのソート済み配列 A1 と A2、および整数 k が与えられているとします。ここでペア (u, v) は「A1 から1つの要素」と「A2 から1つの要素」を組み合わせたものとして定義します。このとき、要素の合計が小さい順に k 個のペア [(u₁, v₁), (u₂, v₂), …, (uₖ, vₖ)] を求めるのが目的です。

たとえば、A1 = [1, 7, 11]、A2 = [2, 4, 6]、k = 3 が入力された場合、出力は [(1, 2), (1, 4), (1, 6)] となります。

解法のアプローチ

この問題は優先度付きキュー（ヒープ）を活用することで効率的に解けます。ポイントは、A1 の各要素についてまず「その要素と A2[0] のペア」だけをキューに入れておき、最小合計のペアを取り出すたびに、同じ A1 の要素における次の候補 (A1[i], A2[j+1]) をキューへ追加していくという発想です。

なお、std::priority_queue はデフォルトでは最大ヒープとして動作するため、Comparator で比較結果を反転させることで最小ヒープとして利用しています。

アルゴリズムの手順

• 2つの値 firstVal・secondVal とインデックス idx を保持する構造体 Data を定義します。
• Data 型を要素とし、合計値で大小を比較する Comparator を持つ priority_queue を作成します。
• n ← A1 のサイズ、m ← A2 のサイズ とします。
• n または m が 0 の場合、空の結果を返します。
• 結果格納用の二次元ベクタ ret を用意します。
• i = 0 ～ n−1 の各 i について、(A1[i], A2[0], 0) をデータとしてキューに挿入します。
• キューが空でなく、かつ k が残っている間、以下を繰り返します。
  • キューの先頭要素 curr を取り出して削除します。
  • (curr.firstVal, curr.secondVal) を ret に追加します。
  • curr.idx + 1 < m である場合、(curr.firstVal, A2[curr.idx + 1], curr.idx + 1) をキューに挿入します。
  • k を 1 減らします。
• 最後に ret を返します。

C++による実装例

以下のコードで実際の動作を確認できます。

#include <bits/stdc++.h>
#include <stack>
using namespace std;
struct Data{
    int firstVal, secondVal, idx;
    Data(int a, int b, int c){
        firstVal = a;
        secondVal = b;
        idx = c;
    }
};
struct Comparator{
    bool operator()(Data a, Data b){
        return !(a.firstVal + a.secondVal < b.firstVal + b.secondVal);
    }
};
class Solution {
    public:
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        priority_queue <Data, vector <Data>, Comparator> pq;
        int n = nums1.size();
        int m = nums2.size();
        if(!n || !m)return {};
        vector < vector <int> > ret;
        for(int i = 0; i < n; i++){
            pq.push(Data(nums1[i], nums2[0], 0));
        }
        while(!pq.empty() && k--){
            Data curr = pq.top();
            pq.pop();
            ret.push_back({curr.firstVal, curr.secondVal});
            if(curr.idx + 1 < m){
                pq.push(Data(curr.firstVal, nums2[curr.idx + 1], curr.idx + 1));
            }
        }
        return ret;
    }
};
void print(vector <int> const &arr) {
    cout<<"[";
    for(int i=0; i < arr.size(); i++)
        std::cout << arr.at(i) <<",";
    cout<<"]";
}
int main() {
    vector<int> nums1{1,7,11};
    vector<int> nums2{2,4,6};
    int k = 3;
    Solution ob1;
    vector<vector<int>> numsRet;
    numsRet = ob1.kSmallestPairs(nums1, nums2, k);
    cout<<"[";
    for (vector<int> x : numsRet) {
        print(x);
        cout<<",";
    }
    cout<<"]"<<endl;
    return 0;
}

入力

nums1 = [1, 7, 11]
nums2 = [2, 4, 6]
k = 3

出力

[[1,2],[1,4],[1,6]]

計算量の目安

初期化段階で n 個の要素をヒープに挿入するのに O(n log n)、その後は最大 k 回のポップとプッシュを行うため O(k log n) が必要です。したがって全体の時間計算量は O((n + k) log(n + k))、空間計算量は O(n + k) となります。全ペアを生成してソートする素朴な手法（O(nm log nm)）と比べ、特に k が小さい場合に大幅な高速化が期待できます。