C++で解説：範囲合計クエリと平方根による範囲更新をBITで高速化する方法

配列と複数のクエリが与えられます。クエリには次の2種類があります。

• update[ L, R ]：L番目からR番目までの各要素を、その平方根の値に更新する
• query[ L, R ]：L番目からR番目までの要素の合計を求める

ここでは、1始まりのインデックスを持つ配列を想定します。以下に入力例と出力例を示します。

入力: nums[ ] = { 0, 9, 4, 1, 5, 2, 3 }, Query[ ] = { {1, 1, 3}, {2, 1, 2}, {1, 2, 5}, { 1, 4, 5}}
出力: 14
10
7

• 1つ目のクエリの最初の要素が「1」なので、1〜3の範囲合計を計算します。すなわち 9 + 4 + 1 = 14
• 2つ目のクエリの最初の要素が「2」なので、1〜2の範囲の要素を平方根に更新します。新しい配列は { 3, 2, 1, 5, 2, 3 } になります
• 3つ目のクエリの最初の要素が「1」なので、2〜5の範囲合計を計算します。すなわち 2 + 1 + 5 + 2 = 10
• 4つ目のクエリの最初の要素が「1」なので、4〜5の範囲合計を計算します。すなわち 5 + 2 = 7

入力: nums[] = { 0, 3, 2, 4, 16, 2 }, Query[ ] = {{1, 1, 3}, {2, 2, 5}}
出力: 9

解法のアプローチ

単純なアプローチ

クエリの数だけループを回し、合計クエリの場合は範囲の合計を返し、更新クエリの場合は配列を直接更新するという方法が考えられます。しかし、この方法の時間計算量は O(q × n) となり、データ規模が大きい場合は非効率です。より効率的なアプローチを見ていきましょう。

効率的なアプローチ

操作の回数や反復処理の回数を減らすことで、プログラムを効率化できます。そこでBinary Indexed Tree（BIT / フェニック木）を使用します。BITでは、内部に配列を持ち、更新と累積和取得のための2つの関数を用意します。

更新クエリについては、要素がすでに「1」であれば、その平方根も1であるため更新する必要がありません。そこで「1より大きい要素のインデックス」を set（平衡二分探索木）に格納しておき、lower_bound による二分探索でL番目以降のインデックスを素早く特定し、範囲内のすべての要素が更新されるまで順に処理を進めます。更新後の値が1になったインデックスは、それ以降どのような更新クエリでも常に1のままであるため、setから削除してしまいます。

この工夫により、各要素が実際に更新される回数はごくわずかに抑えられ（大きな値でも平方根の適用を繰り返せばすぐに1になるため）、全体の計算量はほぼ O((n + q) log n) に収まります。

合計クエリについては、query(R) − query(L−1) を計算することで、範囲 [L, R] の合計を求めることができます。

実装例

上記アプローチのC++コード

#include <bits/stdc++.h>
using namespace std;
// 入力配列の最大サイズ
const int m = 200;
// Binary Indexed Tree（BIT）の作成
int binary_indexed[m + 1];
// 更新クエリ用の関数
void update_q(int a, int x, int n){
    while(a <= n) {
        binary_indexed[a] += x;
        a += a & -a;
    }
}
// 累積和（先頭からの合計）を計算する関数
int sum_q(int a){
    int s = 0;
    while(a > 0) {
        s += binary_indexed[a];
        a -= a & -a;
    }
    return s;
}
int main(){
    int no_query = 4;
    int nums[] = {  0, 9, 4, 1, 5, 2, 3 };
    int n = sizeof(nums) / sizeof(nums[0]);
    // クエリを格納する2次元配列
    int q[no_query + 1][3];
    q[0][0] = 1, q[0][1] = 1, q[0][2] = 3;
    q[1][0] = 2, q[1][1] = 1, q[1][2] = 2;
    q[2][0] = 1, q[2][1] = 2, q[2][2] = 5;
    q[3][0] = 1, q[3][1] = 4, q[3][2] = 5;
    set<int> s;
    for (int i = 1; i < n; i++) {
        // 1より大きい要素のインデックスをsetに挿入
        if (nums[i] > 1)
            s.insert(i);
        update_q(i, nums[i], n);
    }
    for (int i = 0; i < no_query; i++) {
        // 先頭の値で更新クエリか合計クエリかを判定
        if (q[i][0] == 2) {
            while (true) {
                // 二分探索でL以降の最小インデックスを取得
                auto it = s.lower_bound(q[i][1]);
                // 右端を超えたら終了
                if (it == s.end() || *it > q[i][2])
                    break;
                q[i][1] = *it;
                // 要素をその平方根の値に更新
                update_q(*it, (int)sqrt(nums[*it]) - nums[*it], n);
                nums[*it] = (int)sqrt(nums[*it]);
                // 更新後の値が1なら、以降は更新不要のためsetから削除
                if (nums[*it] == 1)
                    s.erase(*it);
                q[i][1]++;
            }
        } else {
            cout <<"query" << i+1 <<": " << (sum_q(q[i][2]) - sum_q(q[i][1] - 1)) << endl;
        }
    }
    return 0;
}

出力結果

query1: 14
query3: 10
query4: 7

まとめ

このチュートリアルでは、配列に対する範囲合計クエリと、要素を平方根へ置き換える範囲更新クエリについて解説しました。まず単純なアプローチの課題（O(q × n) の計算量）を確認し、続いて Binary Indexed Tree（BIT）と set を組み合わせた効率的なアプローチを紹介しました。紹介したC++の実装は、C、Java、Pythonなど他のプログラミング言語にも容易に応用できます。本記事が皆さんのアルゴリズム学習の一助になれば幸いです。