TheAlgorithms-C-Plus-Plus/range_queries/bit.cpp

77 lines
1.3 KiB
C++
Raw Normal View History

2019-10-04 21:25:54 +08:00
// Binary Indexed Tree.
2020-04-18 10:43:43 +08:00
#include <iostream>
2019-10-09 00:01:36 +08:00
2019-10-04 21:25:54 +08:00
using namespace std;
2019-10-09 00:01:36 +08:00
class Bit
{
int n;
2019-10-09 00:02:29 +08:00
vector<int> bit;
2019-10-09 00:01:36 +08:00
inline int offset(int x)
{
return (x & (-x));
}
2019-10-04 21:25:54 +08:00
2019-10-09 00:01:36 +08:00
public:
2019-10-04 21:25:54 +08:00
2019-10-09 00:02:29 +08:00
Bit(vector<int>& arr)
2019-10-09 00:01:36 +08:00
{
n = arr.size();
bit.assign(n + 1, 0);
for (int i = 0; i < n; ++i)
{
update(i, arr[i]);
}
}
Bit(int x)
{
n = x;
bit.assign(n + 1, 0);
}
2019-10-04 21:25:54 +08:00
2019-10-09 00:01:36 +08:00
void update(int id, int val)
{
// Add val at id
id++;
while (id <= n)
{
bit[id] += val;
id += offset(id);
}
2019-10-04 21:25:54 +08:00
}
2019-10-09 00:01:36 +08:00
int sum(int id)
{
// Get prefix sum upto id.
id++;
int res = 0;
while (id > 0)
{
res += bit[id];
id -= offset(id);
}
return res;
2019-10-04 21:25:54 +08:00
}
2019-10-09 00:01:36 +08:00
int sum_range(int l, int r)
{
return sum(r) - sum(l - 1);
}
2019-10-04 21:25:54 +08:00
};
2019-10-09 00:01:36 +08:00
int main()
{
int n = 5;
2019-10-09 00:02:29 +08:00
vector<int> arr = { 1, 2, 3, 4, 5 };
2019-10-09 00:01:36 +08:00
Bit x(arr);
2019-10-04 21:25:54 +08:00
2019-10-09 00:01:36 +08:00
assert(x.sum_range(0, 0) == 1);
assert(x.sum_range(0, 1) == 3);
assert(x.sum_range(0, 2) == 6);
x.update(0, 6);
assert(x.sum_range(0, 0) == 6);
assert(x.sum_range(0, 1) == 8);
assert(x.sum_range(0, 2) == 11);
return 0;
2019-10-04 21:25:54 +08:00
}