C++でマージソートツリーを実装して範囲クエリを高速化する

整数配列、セグメントの開始・終了インデックス、キー値 k が与えられたとき、指定された範囲内で k 番目に小さい要素 を効率的に求める問題を解説します。この問題は「マージソートツリー（Merge Sort Tree）」と呼ばれるデータ構造を用いると、クエリごとに O(log² N) で処理できます。

問題の理解：具体例で確認

例 1

入力配列: arr[] = {7, 8, 1, 4, 6, 8, 10}
クエリ 1: start = 2, end = 4, k = 2  → 範囲 [8, 1, 4] で 2 番目に小さい値
クエリ 2: start = 1, end = 6, k = 3  → 範囲 [7, 8, 1, 4, 6, 8] で 3 番目に小さい値

出力: 4, 6

解説:

• インデックス 2〜4 の部分配列 [8, 1, 4] をソートすると [1, 4, 8] で、2 番目は 4
• インデックス 1〜6 の部分配列 [7, 8, 1, 4, 6, 8] をソートすると [1, 4, 6, 7, 8, 8] で、3 番目は 6

例 2

入力配列: arr[] = {2, 7, 9, 4, 6, 5, 1}
クエリ 1: start = 3, end = 6, k = 4  → 範囲 [9, 4, 6, 5] で 4 番目に小さい値
クエリ 2: start = 2, end = 5, k = 3  → 範囲 [7, 9, 4, 6] で 3 番目に小さい値

出力: 9, 7

マージソートツリーの仕組みとアプローチ

マージソートツリーは、セグメントツリーの各ノードに「担当区間の要素をソートした配列」を持たせた構造です。マージソートのマージ操作を利用して構築します。

構築手順

1. 要素とインデックスのペアを作成: (値, 元のインデックス) のペアをベクターに格納し、値でソートする
2. セグメントツリーを再帰的に構築:
   • 葉ノードでは、担当する単一要素の「元のインデックス」を格納
   • 内部ノードでは、左右の子ノードの配列を std::merge でマージしてソート済み配列を作成
3. 各ノードの配列には「値のソート順での元インデックス」が昇順で並ぶ

クエリ処理（k 番目に小さい要素の検索）

根から葉へ向かって二分探索的に降りていきます。

1. 現在のノードで、左の子ノードの配列内に「クエリ範囲 [queryStart, queryEnd] に含まれるインデックス」がいくつあるかを lower_bound/upper_bound で数える（これを M とする）
2. M >= k なら、答えは左の子ノードの区間にある → 左へ進む
3. M < k なら、答えは右の子ノードの区間にあり、順位は k - M 番目になる → 右へ進む
4. 葉に到達したら、その要素の元インデックスを返す

計算量: 構築 O(N log N)、クエリ O(log² N)

C++ 実装例

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

const int MAX = 1000; // 配列サイズの上限に合わせて調整

// セグメントツリーの構築
// tree[node] には、担当区間の要素を「値でソートしたときの元インデックス」が昇順で入る
void generateTree(int node, int left, int right,
                  const vector<pair<int, int>> &sortedPairs,
                  vector<int> tree[]) {
    if (left == right) {
        tree[node].push_back(sortedPairs[left].second);
        return;
    }
    int mid = (left + right) / 2;
    generateTree(node * 2, left, mid, sortedPairs, tree);
    generateTree(node * 2 + 1, mid + 1, right, sortedPairs, tree);
    merge(tree[node * 2].begin(), tree[node * 2].end(),
          tree[node * 2 + 1].begin(), tree[node * 2 + 1].end(),
          back_inserter(tree[node]));
}

// k 番目に小さい要素のインデックスを再帰的に探索
int calculateKSmallest(int segLeft, int segRight,
                       int queryStart, int queryEnd,
                       int node, int k,
                       const vector<int> tree[]) {
    if (segLeft == segRight) {
        return tree[node][0]; // 葉の要素の元インデックス
    }
    int mid = (segLeft + segRight) / 2;
    
    // 左の子ノードの配列内で、クエリ範囲に含まれるインデックスの個数を数える
    int leftChild = node * 2;
    int lastInRange  = upper_bound(tree[leftChild].begin(), tree[leftChild].end(), queryEnd)   - tree[leftChild].begin();
    int firstInRange = lower_bound(tree[leftChild].begin(), tree[leftChild].end(), queryStart) - tree[leftChild].begin();
    int countInLeft  = lastInRange - firstInRange;
    
    if (countInLeft >= k) {
        // 答えは左側にある
        return calculateKSmallest(segLeft, mid, queryStart, queryEnd, leftChild, k, tree);
    } else {
        // 答えは右側にある（順位を調整）
        return calculateKSmallest(mid + 1, segRight, queryStart, queryEnd,
                                  leftChild + 1, k - countInLeft, tree);
    }
}

// クエリ用ラッパー（1-indexed 入力を 0-indexed に変換）
int queryWrapper(int queryStart, int queryEnd, int k, int n,
                 const vector<pair<int, int>> &sortedPairs,
                 vector<int> tree[]) {
    return calculateKSmallest(0, n - 1, queryStart - 1, queryEnd - 1, 1, k, tree);
}

int main() {
    int input[] = {7, 8, 1, 4, 6, 8, 10};
    int n = sizeof(input) / sizeof(input[0]);
    
    // (値, 元インデックス) のペアを作って値でソート
    vector<pair<int, int>> sortedPairs;
    for (int i = 0; i < n; ++i) {
        sortedPairs.emplace_back(input[i], i);
    }
    sort(sortedPairs.begin(), sortedPairs.end());
    
    // セグメントツリーの配列
    vector<int> tree[MAX];
    generateTree(1, 0, n - 1, sortedPairs, tree);
    
    cout << "範囲内で k 番目に小さい要素: " << endl;
    
    // クエリ 1: インデックス 2〜4 (1-indexed) で 2 番目に小さい
    int idx = queryWrapper(2, 4, 2, n, sortedPairs, tree);
    cout << input[idx] << endl; // 4
    
    // クエリ 2: インデックス 1〜6 で 3 番目に小さい
    idx = queryWrapper(1, 6, 3, n, sortedPairs, tree);
    cout << input[idx] << endl; // 6
    
    return 0;
}

実行結果

範囲内で k 番目に小さい要素:
4
6

ポイントまとめ

• マージソートツリーは「範囲内の k 番目に小さい値」や「範囲内で x 以下の要素数」などのクエリに強い
• 各ノードにソート済み配列を持つため、メモリは O(N log N) 必要
• クエリごとに O(log² N) で回答可能（各レベルで二分探索を行うため）
• 静的な配列に対するオフラインクエリに適している（更新が必要なら Wavelet Tree や BIT + オフライン処理などを検討）

この実装をベースに、要素の値そのものを返すようにしたり、複数クエリを高速に処理したりする応用が可能です。