C++でk個のソート済みリストをカバーする最小範囲を求めるアルゴリズム

この記事では、k個のソートされた整数リストが与えられたときに、各リストから少なくとも1つの数値を含む最小の範囲を求める問題について解説します。

ここで、範囲[a, b]が範囲[c, d]より「小さい」とは、次のいずれかの条件を満たす場合を指します。

• b − a < d − c(幅が狭い)
• b − a == d − c かつ a < c(幅が同じ場合は開始位置が小さい方)

入力と出力の例

たとえば、入力が以下のような3つのリストだったとします。

[[4,10,15,25,26], [0,9,14,20], [5,18,24,30]]

この場合、出力は [14, 18] となります。14は2番目のリスト、18は3番目のリスト、そして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 ループ(無限ループ)で以下を実行する:
  • pq の先頭要素を temp として取得し、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 を返す

この手法では、毎回ヒープから最小値を取り出し、その最小値が属するリストの次の要素をヒープに戻すことで、常に「全リストをカバーする範囲」の中で最も狭い候補を探索できます。計算量は O(n log n)(n は全要素数)程度に抑えられます。

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]

まとめ

k個のソート済みリストから最小の範囲を求める問題は、最小ヒープと各リストへのポインタを組み合わせることで、全要素を高々一度ずつ処理するだけで解決できます。ポイントは、ヒープ内の最小値を基準に範囲を評価し、最小値が属するリストだけを前進させる点です。このテクニックは「k-wayマージ」系の問題にも応用できるので、ぜひ覚えておきましょう。