C++で「右側の区間」を見つけるアルゴリズムを解説

問題概要

区間のリストが与えられたとき、各区間 i について、「始点が区間 i の終点以上であるような区間 j」が存在するかどうかを調べます。このとき、区間 j は区間 i の「右側」にあると言います。

各区間 i に対して、条件を満たす区間 j のうち始点が最小となるもののインデックスを記録します。該当する区間が存在しない場合は -1 を格納し、最終的に各区間に対応する値を配列として出力します。

たとえば、入力が [[3,4], [2,3], [1,2]] の場合、出力は [-1, 0, 1] となります。

• [3, 4]: 右側に位置する区間が存在しないため -1
• [2, 3]: 始点が 3 以上の区間のうち最小のものは [3, 4](インデックス 0)
• [1, 2]: 始点が 2 以上の区間のうち最小のものは [2, 3](インデックス 1)

解法のアプローチ

この問題は、std::map(平衡二分探索木)と lower_bound を組み合わせることで効率的に解けます。手順は以下の通りです。

• n を区間配列のサイズとし、サイズ n の配列 ret を作成してすべて -1 で初期化します。さらに、マップ m を用意します。
• i を 0 から区間の個数までループさせます。
  • intervals[i][0](始点)がすでに m に登録されている場合はスキップします。
  • m[intervals[i][0]] = i + 1 として、始点をキー、インデックス + 1 を値として登録します。
• i を n - 1 から 0 まで逆順にループさせます。
  • it を、intervals[i][1](終点)以上のキーのうち最小のものを指すイテレータ(lower_bound の結果)とします。
  • 該当する要素が存在しない場合(end() を指す場合)は次の反復へ進みます。
  • ret[i] = it の値 − 1 とします。
• ret を返します。

値を「インデックス + 1」として登録しているのは、有効な要素との区別を明確につけるためです。取り出す際に 1 を引けば元のインデックスが復元できます。

計算量は、マップへの挿入と検索がそれぞれ O(log n) であることから、全体で時間計算量 O(n log n)、空間計算量 O(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;
}

class Solution {
public:
    vector<int> findRightInterval(vector<vector<int>>& intervals) {
        int n = intervals.size();
        vector<int> ret(n, -1);
        map<int, int> m;
        // 始点をキー、インデックス+1を値としてマップに登録
        for(int i = 0; i < intervals.size(); i++){
            if(m.count(intervals[i][0])) continue;
            m[intervals[i][0]] = i + 1;
        }
        // 各区間の終点以上の最小の始点を二分探索で取得
        for(int i = n - 1; i >= 0; i--){
            map<int, int>::iterator it = m.lower_bound(intervals[i][1]);
            if(it == m.end()) continue;
            ret[i] = it->second - 1;
        }
        return ret;
    }
};

main(){
    vector<vector<int>> v = {{3,4},{2,3},{1,2}};
    Solution ob;
    print_vector(ob.findRightInterval(v));
}

入力

[[3,4],[2,3],[1,2]]

出力

[-1, 0, 1]