2.4.6 Segment Tree (Range Update)

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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255

/*

Maintain a fixed-size array while supporting range updates and range aggregate queries. Lazy
propagation lets a segment tree defer work below intervals covered by an update: each node stores an
aggregate value and a pending update that is pushed to its children only when needed.

The query operation is defined by an associative aggregate function `combine(a, b)`. The default
code below assumes a numerical array type, defining queries for the "min" of the target range.
Another possible query operation is "sum", in which case `combine(a, b)` should return `a + b`.

Range updates are defined by `apply_delta(v, d, len)`, which applies an update delta `d` to an
aggregate summary `v` representing `len` array values, and by `compose_deltas(old, d)`, which
combines a pending older delta with a newer delta in that order. These functions do not support
arbitrary combinations: applying a delta to a combined segment must be equivalent to applying it to
each child segment and then combining the results, and composed deltas must be equivalent to
performing their updates sequentially. The default code below defines range assignment. For range
increment, `compose_deltas(old, d)` should return `old + d`; `apply_delta(v, d, len)` should return
`v + d` for range-min/range-max queries, and `v + d * len` for range-sum queries.

- `LazySegTree<T>(n, v = T{})` constructs an array of size `n` with indices $[0, `n`)$, and all
  values initialized to `v`.
- `LazySegTree<T>(lo, hi)` constructs an array from the half-open random-access iterator range
  $[`lo`, `hi`)$.
- `size()` returns the size of the array.
- `at(i)` returns the value at index `i`.
- `query(lo, hi)` returns the aggregate of the values at indices in $[`lo`, `hi`]$. If `lo == hi`,
  then the single specified value is returned.
- `update(i, d)` modifies the value at index `i` by applying the delta `d`.
- `update(lo, hi, d)` modifies the value at each array index in $[`lo`, `hi`]$ by applying the delta
  `d` to each value.
- `max_right(lo, pred)` returns the largest boundary `hi` such that the aggregate over the half-open
  range $[`lo`, `hi`)$ satisfies `pred()`. As `hi` increases, `pred()` applied to this aggregate may
  change only from true to false. The empty range is valid, and `size()` is returned if the
  predicate remains true to the end.
- `min_left(hi, pred)` returns the smallest boundary `lo` such that the aggregate over the half-open
  range $[`lo`, `hi`)$ satisfies `pred()`. As `lo` decreases, `pred()` applied to this aggregate may
  change only from true to false. The empty range is valid, and $0$ is returned if the predicate
  remains true to the beginning.

For the boundary-search functions, `pred()` takes aggregate `T` values. For `combine = min`,
`pred(mn) = (mn > x)` makes `max_right()` stop at the first value `<= x` when extending right, while
`min_left()` stops just after the first such value when extending left. With nonnegative values and
`combine = sum`, `pred(sum) = (sum <= x)` finds the longest extension with sum at most `x`.

Time Complexity:
- O(n) per call to both constructors, where $n$ is the size of the array.
- O(1) per call to `size()`.
- O(log n) per call to `at()`, `query()`, `update()`, `max_right()`, and `min_left()`.

Space Complexity:
- O(n) for storage of the array elements.
- O(log n) auxiliary stack space for `query()`, `update()`, `max_right()`, and `min_left()`.
- O(1) auxiliary for `size()`.

*/

#include <algorithm>
#include <cassert>
#include <cstdint>
#include <optional>
#include <vector>

template<typename T>
class LazySegTree {
  static T combine(const T &a, const T &b) { return std::min(a, b); }
  static T apply_delta(const T &v, const T &d, int64_t len) { return d; }
  static T compose_deltas(const T &d1, const T &d2) { return d2; }

  int len;
  std::vector<T> value, delta;
  std::vector<char> pending;

  template<typename Gen>
  void build(int i, int lo, int hi, const Gen &gen) {
    if (lo == hi) {
      value[i] = gen(lo);
      return;
    }
    int mid = lo + (hi - lo) / 2;
    build(i * 2 + 1, lo, mid, gen);
    build(i * 2 + 2, mid + 1, hi, gen);
    value[i] = combine(value[i * 2 + 1], value[i * 2 + 2]);
  }

  void push_delta(int i, int lo, int hi) {
    if (pending[i]) {
      value[i] = apply_delta(value[i], delta[i], hi - lo + 1);
      if (lo != hi) {
        int l = i * 2 + 1, r = i * 2 + 2;
        delta[l] = pending[l] ? compose_deltas(delta[l], delta[i]) : delta[i];
        delta[r] = pending[r] ? compose_deltas(delta[r], delta[i]) : delta[i];
        pending[l] = pending[r] = true;
      }
      pending[i] = false;
    }
  }

  T query(int i, int lo, int hi, int tgt_lo, int tgt_hi) {
    push_delta(i, lo, hi);
    if (lo == tgt_lo && hi == tgt_hi) {
      return value[i];
    }
    int mid = lo + (hi - lo) / 2;
    if (tgt_lo <= mid && mid < tgt_hi) {
      return combine(
          query(i * 2 + 1, lo, mid, tgt_lo, std::min(tgt_hi, mid)),
          query(i * 2 + 2, mid + 1, hi, std::max(tgt_lo, mid + 1), tgt_hi)
      );
    }
    if (tgt_lo <= mid) {
      return query(i * 2 + 1, lo, mid, tgt_lo, std::min(tgt_hi, mid));
    }
    return query(i * 2 + 2, mid + 1, hi, std::max(tgt_lo, mid + 1), tgt_hi);
  }

  void update(int i, int lo, int hi, int tgt_lo, int tgt_hi, const T &d) {
    push_delta(i, lo, hi);
    if (hi < tgt_lo || lo > tgt_hi) {
      return;
    }
    if (tgt_lo <= lo && hi <= tgt_hi) {
      delta[i] = d;
      pending[i] = true;
      push_delta(i, lo, hi);
      return;
    }
    update(i * 2 + 1, lo, lo + (hi - lo) / 2, tgt_lo, tgt_hi, d);
    update(i * 2 + 2, lo + (hi - lo) / 2 + 1, hi, tgt_lo, tgt_hi, d);
    value[i] = combine(value[i * 2 + 1], value[i * 2 + 2]);
  }

  template<typename Pred>
  int max_right(int i, int lo, int hi, int tgt_lo, const Pred &pred, std::optional<T> &acc) {
    push_delta(i, lo, hi);
    if (hi < tgt_lo) {
      return -1;
    }
    if (tgt_lo <= lo) {
      T next = acc ? combine(*acc, value[i]) : value[i];
      if (pred(next)) {
        acc = next;
        return -1;
      }
      if (lo == hi) {
        return lo;
      }
    }
    int mid = lo + (hi - lo) / 2;
    int res = max_right(i * 2 + 1, lo, mid, tgt_lo, pred, acc);
    return res != -1 ? res : max_right(i * 2 + 2, mid + 1, hi, tgt_lo, pred, acc);
  }

  template<typename Pred>
  int min_left(int i, int lo, int hi, int tgt_hi, const Pred &pred, std::optional<T> &acc) {
    push_delta(i, lo, hi);
    if (tgt_hi <= lo) {
      return -1;
    }
    if (hi < tgt_hi) {
      T next = acc ? combine(value[i], *acc) : value[i];
      if (pred(next)) {
        acc = next;
        return -1;
      }
      if (lo == hi) {
        return lo + 1;
      }
    }
    int mid = lo + (hi - lo) / 2;
    int res = min_left(i * 2 + 2, mid + 1, hi, tgt_hi, pred, acc);
    return res != -1 ? res : min_left(i * 2 + 1, lo, mid, tgt_hi, pred, acc);
  }

 public:
  explicit LazySegTree(int n, const T &v = T{}) : len(n), pending(4 * len, false) {
    assert(len > 0);
    value.assign(4 * len, v);
    delta.assign(4 * len, v);
    build(0, 0, len - 1, [&](int) { return v; });
  }

  template<typename It>
  LazySegTree(It lo, It hi) : len(hi - lo), pending(4 * len, false) {
    assert(len > 0);
    value.assign(4 * len, *lo);
    delta.assign(4 * len, *lo);
    build(0, 0, len - 1, [&](int i) { return *(lo + i); });
  }

  int size() const { return len; }

  T at(int i) {
    assert(0 <= i && i < len);
    return query(i, i);
  }

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

  void update(int i, const T &d) {
    assert(0 <= i && i < len);
    update(0, 0, len - 1, i, i, d);
  }

  void update(int lo, int hi, const T &d) {
    assert(0 <= lo && lo <= hi && hi < len);
    update(0, 0, len - 1, lo, hi, d);
  }

  template<typename Pred>
  int max_right(int lo, const Pred &pred) {
    assert(0 <= lo && lo <= len);
    std::optional<T> acc;
    int res = max_right(0, 0, len - 1, lo, pred, acc);
    return res == -1 ? len : res;
  }

  template<typename Pred>
  int min_left(int hi, const Pred &pred) {
    assert(0 <= hi && hi <= len);
    std::optional<T> acc;
    int res = min_left(0, 0, len - 1, hi, pred, acc);
    return res == -1 ? 0 : res;
  }
};

/*** Example Usage ***/

using namespace std;

int main() {
  vector<int> a{6, -2, 1, 8, 10};
  LazySegTree<int> t(a.begin(), a.end());
  t.update(2, 4);
  vector<int> expected{6, -2, 4, 8, 10};
  for (int i = 0; i < t.size(); i++) {
    assert(t.at(i) == expected[i]);
  }
  assert(t.query(0, 3) == -2);
  t.update(0, 4, 5);
  t.update(3, 2);
  t.update(3, 1);
  expected = {5, 5, 5, 1, 5};
  for (int i = 0; i < t.size(); i++) {
    assert(t.at(i) == expected[i]);
  }
  assert(t.query(0, 3) == 1);

  // Boundary search works through lazy updates too; values are now {5, 5, 5, 1, 5}.
  assert(t.max_right(0, [](int mn) { return mn > 1; }) == 3);
  assert(t.min_left(5, [](int mn) { return mn > 1; }) == 4);
  return 0;
}