何回かに分けてセグメント木ライブラリを整備します。今回は、1点更新と範囲の和を求めることに集中します。
セグメント木(Range Sum Query 固定版)
前回は、セグメント木を学び、1点更新と範囲の最小値を求めるプログラムを紹介しました。今回は、1点更新と範囲の和を求めることに固定したセグメント木の実装を紹介します。
コードは以下となります。
ライブラリのテスト
AIZU ONLINE JUDGE で出題されている問題でアルゴリズムをテストします。以下が問題のリンク先です。
問題の注意点としては、クエリ0(add)は、与えられた x を加えることと、数列の要素が1から始まっていることです。DSL_2_A問題は、数列の要素は0から始まっていました。
この問題を解くプログラムは以下です。SegmentTree クラスで前回からコードを変更した箇所は背景色を変更しています。AC判定となります。
#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
class SegmentTree {
private:
vector<int> data;
public:
int size;
SegmentTree(int n) {
size = 1;
while (size < n) {
size *= 2;
}
data.assign(2 * size - 1, 0);
}
void update(int k, int a) {
k += size - 1;
data[k] += a;
while (k > 0) {
k = (k - 1) / 2;
data[k] = data[k * 2 + 1] + data[k * 2 + 2];
}
}
int query(int a, int b) {
return query_sub(a, b, 0, 0, size);
}
// call ST.query(a, b, 0, 0, size);
int query_sub(int a, int b, int k, int l, int r) {
if ((r <= a)||(b <= l)) {
return 0;
}
if ((a <= l)&&(r <= b)) {
return data[k];
} else {
int vl = query_sub(a, b, k * 2 + 1, l, (l + r) / 2);
int vr = query_sub(a, b, k * 2 + 2, (l + r) / 2, r);
return vl + vr;
}
}
};
int main()
{
int n, q;
cin >> n >> q;
SegmentTree st(n);
for (int i = 0; i < q; ++i) {
int command;
cin >> command;
if (command == 0) {
int pos, x;
cin >> pos >> x;
--pos;
st.update(pos, x);
} else if (command == 1) {
int L, R;
cin >> L >> R;
--L;
cout << st.query(L, R) << endl;
}
}
return 0;
}
最後に
2つの近い問題を解くことにより、セグメント木のライブラリを抽象化できそうなことに気づきます。次回は、汎用化したライブラリを紹介します。