2.6.1 Fenwick Tree (Point Update, Range 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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

/*

Maintain a numerical array while supporting point increments and range-sum queries. A Fenwick tree
(also called a binary indexed tree) stores partial prefix sums: each internal entry `tree[i]` holds
the sum of a block of indices whose length is the lowest set bit of `i`. A point update touches
O(log n) covering blocks, and a prefix-sum query combines O(log n) disjoint blocks.

When prefix sums are nondecreasing, `max_prefix(c)` answers prefix-boundary queries by binary
lifting: it walks the implicit power-of-two block structure to find the longest prefix with sum at
most `c`. This lets the same structure double as a dynamic multiset over a bounded integer domain,
useful for coordinate-compressed $k$-th element lookups and online rank queries.

- `Fenwick<T>(n)` constructs an array with 0-based indices $[0, `n`)$, with all values initialized
  to $0$.
- `size()` returns the size of the array.
- `add(i, x)` adds `x` to the value at index `i`.
- `set(i, x)` assigns the value at index `i` to `x`.
- `at(i)` returns the value at index `i`.
- `sum(hi)` returns the sum of all values at indices $[0, `hi`]$.
- `sum(lo, hi)` returns the sum of all values at indices $[`lo`, `hi`]$.
- `max_prefix(c)` returns the largest boundary `hi` such that `sum(0, hi - 1)` $\leq$ `c`, assuming
  prefix sums are nondecreasing and `c` is nonnegative. It may return any boundary in $[0, `n`]$.

The value type `T` must represent $0$ and support addition and subtraction.

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 `add()`, both `sum()` functions, `set()`, `at()`, and `max_prefix()`.

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

*/

#include <cassert>
#include <vector>

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

 public:
  explicit Fenwick(int n) : len(n), tree(n + 1) {}

  int size() const { return len; }

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

  void set(int i, const T &x) { add(i, x - at(i)); }
  T at(int i) const { return sum(i, i); }

  T sum(int hi) const {
    assert(-1 <= hi && hi < len);
    T res = 0;
    for (++hi; hi > 0; hi -= hi & -hi) {
      res += tree[hi];
    }
    return res;
  }

  T sum(int lo, int hi) const {
    assert(0 <= lo && lo <= hi && hi < len);
    return sum(hi) - sum(lo - 1);
  }

  int max_prefix(T c) const {
    int pos = 0, pw = 1;
    while (pw * 2 <= len) {
      pw *= 2;
    }
    for (; pw > 0; pw >>= 1) {
      if (pos + pw <= len && tree[pos + pw] <= c) {
        c -= tree[pos + pw];
        pos += pw;
      }
    }
    return pos;
  }
};

/*** Example Usage ***/

using namespace std;

int main() {
  vector<int> a{10, 1, 2, 3, 4};
  Fenwick<int> t(5);
  for (int i = 0; i < static_cast<int>(a.size()); i++) {
    t.set(i, a[i]);
  }
  t.add(0, -5);  // a is now {5, 1, 2, 3, 4}.
  assert(t.at(0) == 5);
  assert(t.sum(1, 3) == 6);  // 1 + 2 + 3
  assert(t.sum(4) == 15);    // 5 + 1 + 2 + 3 + 4

  // With nonnegative values, the tree doubles as a frequency table for order-statistic queries.
  Fenwick<int> freq(8);
  freq.add(1, 1);                   // One element of value 1.
  freq.add(3, 3);                   // Three elements of value 3.
  freq.add(6, 1);                   // One element of value 6.
  assert(freq.max_prefix(0) == 1);  // The longest prefix with count at most 0 is [0, 1).
  assert(freq.max_prefix(3) == 3);  // Prefix [0, 3) has count 1; [0, 4) has count 4.
  assert(freq.max_prefix(4) == 6);  // Prefix [0, 6) has count 4; [0, 7) has count 5.
  return 0;
}