Segment Tree
Segment Tree is a complete binary tree which can calculate partial sum at O(log n) by having the calculation results in each partials as node. In the following example, the calculated value is sum, but if it is the minimum value, Range Minimum Query (RMQ) can be solved and if it is the sorted list, merge sort is processed. When updating the value, recalculate in order from the bottom.
#include <bits/stdc++.h>
using namespace std;
template <typename S, typename T>
class SegmentTree {
S default_value;
int leaf_num = 1;
vector<S> nodes;
/*
[0]
/ \
[1] [2]
/ \ / \
[3][4][5][6]
*/
function<int(int, int)> calc;
public:
SegmentTree(vector<S> data, S default_value, function<T(S, S)> calc) {
while (leaf_num < data.size()) leaf_num *= 2;
this->calc = calc;
this->default_value = default_value;
nodes = vector<S>(2 * data.size() - 1, default_value);
for (int i = 0; i < data.size(); i++) update(i, data[i]);
}
void update(int idx, S value) {
idx += leaf_num - 1;
nodes[idx] = value;
while (idx > 0) {
idx = (idx - 1) / 2; // parent
nodes[idx] = this->calc(nodes[idx * 2 + 1], nodes[idx * 2 + 2]);
}
}
T query(int from, int to, int node_idx=0, int left=0, int right=-1) {
if (right == -1) right = leaf_num-1;
if (right < from || to < left) return this->calc(default_value, default_value);
if (from <= left && right <= to) return this->calc(default_value, nodes[node_idx]);
return this->calc(
query(from, to, node_idx * 2 + 1, left, (left + right) / 2), // left child
query(from, to, node_idx * 2 + 2, (left + right) / 2 + 1, right) // right child
);
}
};
int main() {
SegmentTree<int, int> segtree({6, 4, 2, 3, 5, 1, 3, 2}, 0, [](int a, int b) { return a + b; });
cout << segtree.query(0, 6) << endl; // 24
segtree.update(0, 8);
cout << segtree.query(0, 6) << endl; // 26
}
When updating range, if the length is m, it will be O(m log n), but it can be improved by storing the updating value of the range with the range in another Segment Tree and doing lazy evaluation.
Binary Indexed Tree (BIT)
When calculating the partial sum with Segment Tree, the value of the child node on the right is included in the parent node so it can be omitted. Binary Indexed Tree is a tree indexed to the remaining nodes as follows, and the value can be updated or the sum from 0 to idx by adding or subtracting the last 1 bit in the binary representated index. Although it is limited in what it can do compared to Segment Tree, it can be implemented more simply and faster.
The last 1 bit can be obtained with x & -x.
#include <bits/stdc++.h>
using namespace std;
class BinaryIndexedTree {
vector<int> nodes;
/*
[4]
/ \
[2] -
/ \ / \
[1] - [3] -
*/
public:
BinaryIndexedTree(vector<int> data) {
this->nodes = vector<int>(data.size()+1, 0);
for(int i = 0; i < data.size(); i++) this->add(i, data[i]);
}
void add(int idx, int value) {
idx++;
while (idx < nodes.size()) {
nodes[idx] += value;
idx += idx & -idx; // -----> e.g. 0001 & 1001 = 0001
}
}
int sum(int to) {
to++;
int ret = 0;
while (to > 0) {
ret += nodes[to];
to -= to & -to; // - - ->
}
return ret;
}
};
int main() {
BinaryIndexedTree bit({6, 4, 2, 3, 5, 1, 3, 2});
cout << bit.sum(6) << endl; // 24
bit.add(0, 2);
cout << bit.sum(6) << endl; // 26
}