C++でk個のソート済みリストから要素を含む最小範囲を検索する方法

問題概要

k個の異なるリストがあり、それぞれの要素は昇順にソートされているとします。このとき、k個のリストすべてから少なくとも1つの数値を含む最小の範囲 [a, b] を求めます。

範囲の大小関係は次のように定義されます。

• b - a < d - c の場合、範囲 [a, b] は範囲 [c, d] より「小さい」とみなす
• b - a == d - c の場合は、開始位置を比較し、a < c なら [a, b] の方が「小さい」とみなす

例えば、入力が [[4,10,15,25,26], [0,9,14,20], [5,18,24,30]] の場合、出力は [14, 18] となります。14は2番目のリストから、18は3番目のリストから取られ、1番目のリストには15または25が該当します。

解法のアプローチ：優先度付きキュー（ヒープ）

この問題は、各リストの現在位置を指すポインタと優先度付きキュー（最小ヒープ）を組み合わせることで効率的に解けます。各ステップでヒープ内の最小値を取り出し、その要素が属するリストの次の要素に進めることで、常に「全リストから1つずつ要素を含む範囲」を追跡し、その中で最も幅が狭いものを記録していきます。

アルゴリズムの手順

• minRange := ∞、maxRange := -∞、rangeSize := ∞、tempMinRange := ∞、tempMaxRange := -∞ として初期化する
• n := nums のサイズ（リストの個数）
• サイズ n の配列 pointers を定義し、各リストの現在位置を管理する
• 優先度付きキュー pq を作成する
• i := 0 から i < n の間、i を1ずつ増やしながら繰り返す：
  • {nums[i][0], i} を pq に挿入する
  • tempMaxRange := tempMaxRange と nums[i][0] の最大値
• while ループ（無限ループ）で以下を実行する：
  • ペア temp := pq の先頭要素を取得
  • pq から先頭要素を削除
  • tempMinRange := temp.first（現在の最小値）
  • idx := temp.second（その要素が属するリストのインデックス）
  • もし tempMaxRange - tempMinRange < rangeSize ならば：
    • rangeSize := tempMaxRange - tempMinRange
    • minRange := tempMinRange
    • maxRange := tempMaxRange
  • pointers[idx] を1増やす
  • もし pointers[idx] が nums[idx] のサイズと等しくなったら、ループを抜ける
  • そうでなければ：
    • tempMaxRange := tempMaxRange と nums[idx][pointers[idx]] の最大値
    • {nums[idx][pointers[idx]], idx} を pq に挿入する
• サイズ2の配列 ans を定義する
• ans[0] := minRange、ans[1] := maxRange
• ans を返す

いずれかのリストの末尾を超えた時点で探索を終了します。これは、そのリストからこれ以上新しい候補を追加できなくなり、より良い範囲が見つかる可能性がないためです。

C++での実装例

それでは、上記のアルゴリズムを実際のコードで確認してみましょう。

#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<auto> v){
    cout << "[";
    for(int i = 0; i<v.size(); i++){
        cout << v[i] << ", ";
    }
    cout << "]"<<endl;
}
struct Comparator{
    bool operator() (pair <int, int> a, pair <int, int> b){
        return !(a.first < b.first);
    }
};
class Solution {
public:
    vector<int> smallestRange(vector<vector<int>>& nums) {
        int minRange = INT_MAX;
        int maxRange = INT_MIN;
        int rangeSize = INT_MAX;
        int tempMinRange, tempMaxRange, tempRangeSize;
        tempMinRange = INT_MAX;
        tempMaxRange = INT_MIN;
        int n = nums.size();
        vector <int> pointers(n);
        priority_queue < pair <int, int>, vector < pair <int, int> >, Comparator > pq;
        for(int i = 0; i < n; i++){
            pq.push({nums[i][0], i});
            tempMaxRange = max(tempMaxRange, nums[i][0]);
        }
        while(1){
            pair <int, int> temp = pq.top();
            pq.pop();
            tempMinRange = temp.first;
            int idx = temp.second;
            if(tempMaxRange - tempMinRange < rangeSize){
                rangeSize = tempMaxRange - tempMinRange;
                minRange = tempMinRange;
                maxRange = tempMaxRange;
            }
            pointers[idx]++;
            if(pointers[idx] == nums[idx].size())break;
            else{
                tempMaxRange = max(tempMaxRange,
                nums[idx][pointers[idx]]);
                pq.push({nums[idx][pointers[idx]], idx});
            }
        }
        vector <int> ans(2);
        ans[0] = minRange;
        ans[1] = maxRange;
        return ans;
    }
};
main(){
    Solution ob;
    vector<vector<int>> v =
    {{4,10,15,25,26},{0,9,14,20},{5,18,24,30}};
    print_vector(ob.smallestRange(v));
}

入力

{{4,10,15,25,26},{0,9,14,20},{5,18,24,30}};

出力

[14, 18]

計算量について

このアルゴリズムの時間計算量は O(n × log n) です（n は全リストの要素数の総数）。各要素は高々1回ヒープに挿入され、1回取り出されるだけなので効率的です。空間計算量は、ヒープに常時リストの個数分の要素しか保持しないため O(k)（k はリストの個数）で抑えられます。