2.6.2 Fenwick Tree (Range Update, Point Query)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76

/*

Maintain a numerical array while supporting range increments and point queries. This is done by
storing the difference array in a Fenwick tree: adding `x` to [`lo`, `hi`] adds `+x` at `lo` and
`-x` after `hi`, and the value at index `i` is the prefix sum of those differences through `i`.

- `FenwickRUPQ<T>(n)` constructs an array with 0-based indices [0, `n`), with values set to 0.
- `size()` returns the size of the array.
- `at(i)` returns the value at index `i`.
- `add(i, x)` adds `x` to the value at index `i`.
- `add(lo, hi, x)` adds `x` to the values at all indices from `lo` to `hi`, inclusive.

Time Complexity:
- O(n) per call to the constructor, where $n$ is the size of the array.
- O(1) per call to `size()`.
- O(log n) per call to `at()` and both `add()` functions.

Space Complexity:
- O(n) for storage of the array elements.
- O(1) auxiliary for all operations.

*/

#include <vector>

template<class T>
class FenwickRUPQ {
  int len;
  std::vector<T> tree;

 public:
  explicit FenwickRUPQ(int n) : len(n), tree(n + 2) {}

  int size() const { return len; }

  T at(int i) const {
    T res = 0;
    for (i++; i > 0; i -= i & -i) {
      res += tree[i];
    }
    return res;
  }

  void add(int i, const T &x) {
    for (i++; i <= len + 1; i += i & -i) {
      tree[i] += x;
    }
  }

  void add(int lo, int hi, const T &x) {
    add(lo, x);
    add(hi + 1, -x);
  }
};

/*** Example Usage and Output:

Values: 5 10 15 10 10

***/

#include <iostream>
using namespace std;

int main() {
  FenwickRUPQ<int> t(5);
  t.add(0, 1, 5);
  t.add(1, 2, 5);
  t.add(2, 4, 10);
  cout << "Values: ";
  for (int i = 0; i < t.size(); i++) {
    cout << t.at(i) << " ";
  }
  cout << endl;
  return 0;
}