2.4.9 Segment Tree Beats

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

/*

Maintain a fixed-size array under the range "chmin" update `a[i] = min(a[i], t)` for all `i` in a
range, while answering range-sum and range-maximum queries. An ordinary lazy segment tree cannot
represent a clamp like this with a single composable tag, because clamping affects only the entries
that currently exceed `t`. Segment tree beats (a.k.a. the Ji Driver tree) resolves this by storing,
for each node, the strict maximum, the second-largest distinct value, and a count of entries equal
to the maximum. A clamp then has three cases at each node: if `t` $\geq$ `max1`, the clamp changes
nothing and the recursion stops (the prune case); if `max2` < `t` < `max1`, every entry equal to
`max1` drops to `t` and all others are untouched, so the node updates in O(1), the sum falls by
$(`max1` - `t`) * `count`$, and `max1` becomes `t` (the break, or tag, case); otherwise (`t` $\leq$
`max2`) the node is too coarse to update directly, so the clamp recurses into both children and
re-pulls. The lazy tag pushes down to a child only when the child's maximum is larger.

The break condition is what makes this efficient: a potential-function argument on the number of
distinct values along root-to-leaf paths shows the total work over any sequence of clamps is
amortized O(log^2 n) each, even though a single clamp can momentarily descend deep. The same scheme
extends to range "chmax" by tracking the symmetric minimum trio, and to mixed clamp-and-add
workloads by carrying an additional add tag; only the chmin / sum / max core is shown here. Unlike
the generic lazy segment tree, this structure cannot be reduced to a pluggable `combine` and
`apply`, because a clamp acts on only the maximal entries of a node rather than uniformly on all of
them, so each operation set must be hand-written and its amortized bound argued separately.

- `SegTreeBeats<T>(n, v = T{})` constructs an array of size `n`, indices $[0, `n`)$, all equal to
  `v`.
- `SegTreeBeats<T>(lo, hi)` constructs an array from the half-open random-access iterator range
  $[`lo`, `hi`)$.
- `size()` returns the size of the array.
- `chmin(lo, hi, t)` replaces `a[i]` with `min(a[i], t)` for every `i` in $[`lo`, `hi`]$.
- `query_sum(lo, hi)` returns the sum of `a[i]` over `i` in $[`lo`, `hi`]$.
- `query_max(lo, hi)` returns the maximum of `a[i]` over `i` in $[`lo`, `hi`]$.
- `at(i)` returns the value at index `i`.

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^2 n) amortized per call to `chmin()`.
- O(log n) per call to `query_sum()`, `query_max()`, and `at()`.

Space Complexity:
- O(n) for storage of the tree.
- O(log n) auxiliary stack space per `chmin()` and query call.

*/

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

template<typename T>
class SegTreeBeats {
  static constexpr T NEG_INF = std::numeric_limits<T>::lowest();

  int len;
  std::vector<T> max1, max2, sum;
  std::vector<int> cnt;  // Number of entries in the node equal to max1.

  void pull(int i) {
    int l = i * 2 + 1, r = i * 2 + 2;
    sum[i] = sum[l] + sum[r];
    if (max1[l] == max1[r]) {
      max1[i] = max1[l];
      cnt[i] = cnt[l] + cnt[r];
      max2[i] = std::max(max2[l], max2[r]);
    } else if (max1[l] > max1[r]) {
      max1[i] = max1[l];
      cnt[i] = cnt[l];
      max2[i] = std::max(max2[l], max1[r]);
    } else {
      max1[i] = max1[r];
      cnt[i] = cnt[r];
      max2[i] = std::max(max1[l], max2[r]);
    }
  }

  // Lower the node's maximum to t. Valid only when max2 < t < max1, so exactly the cnt entries
  // equal to max1 change and the second maximum is left untouched.
  void apply_chmin(int i, const T &t) {
    sum[i] -= (max1[i] - t) * static_cast<T>(cnt[i]);  // Overflow warning.
    max1[i] = t;
  }

  void push(int i) {
    for (int c : {i * 2 + 1, i * 2 + 2}) {
      if (max1[c] > max1[i]) {
        apply_chmin(c, max1[i]);
      }
    }
  }

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

  void chmin(int i, int lo, int hi, int tgt_lo, int tgt_hi, const T &t) {
    if (tgt_hi < lo || hi < tgt_lo || max1[i] <= t) {
      return;  // Prune: outside the range, or every entry is already at most t.
    }
    if (tgt_lo <= lo && hi <= tgt_hi && (lo == hi || max2[i] < t)) {
      apply_chmin(i, t);  // Break: only the maximal entries are affected.
      return;
    }
    push(i);
    int mid = lo + (hi - lo) / 2;
    chmin(i * 2 + 1, lo, mid, tgt_lo, tgt_hi, t);
    chmin(i * 2 + 2, mid + 1, hi, tgt_lo, tgt_hi, t);
    pull(i);
  }

  T query_sum(int i, int lo, int hi, int tgt_lo, int tgt_hi) {
    if (tgt_hi < lo || hi < tgt_lo) {
      return 0;
    }
    if (tgt_lo <= lo && hi <= tgt_hi) {
      return sum[i];
    }
    push(i);
    int mid = lo + (hi - lo) / 2;
    return query_sum(i * 2 + 1, lo, mid, tgt_lo, tgt_hi) +
           query_sum(i * 2 + 2, mid + 1, hi, tgt_lo, tgt_hi);
  }

  T query_max(int i, int lo, int hi, int tgt_lo, int tgt_hi) {
    if (tgt_hi < lo || hi < tgt_lo) {
      return NEG_INF;
    }
    if (tgt_lo <= lo && hi <= tgt_hi) {
      return max1[i];
    }
    push(i);
    int mid = lo + (hi - lo) / 2;
    return std::max(
        query_max(i * 2 + 1, lo, mid, tgt_lo, tgt_hi),
        query_max(i * 2 + 2, mid + 1, hi, tgt_lo, tgt_hi)
    );
  }

 public:
  explicit SegTreeBeats(int n, const T &v = T{})
      : len(n), max1(4 * n), max2(4 * n), sum(4 * n), cnt(4 * n) {
    assert(len > 0);
    build(0, 0, len - 1, [&](int) { return v; });
  }

  template<typename It>
  SegTreeBeats(It lo, It hi)
      : len(static_cast<int>(hi - lo)), max1(4 * len), max2(4 * len), sum(4 * len), cnt(4 * len) {
    assert(len > 0);
    build(0, 0, len - 1, [&](int i) { return *(lo + i); });
  }

  int size() const { return len; }

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

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

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

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

/*** Example Usage ***/

using namespace std;

int main() {
  vector<int64_t> a{5, 4, 3, 2, 1};
  SegTreeBeats<int64_t> t(a.begin(), a.end());
  assert(t.query_sum(0, 4) == 15);
  assert(t.query_max(0, 4) == 5);

  t.chmin(0, 2, 3);  // {3, 3, 3, 2, 1}: only the 5 and 4 are clamped down to 3.
  assert(t.at(0) == 3);
  assert(t.at(1) == 3);
  assert(t.query_sum(0, 4) == 12);
  assert(t.query_max(0, 4) == 3);

  t.chmin(0, 4, 2);  // {2, 2, 2, 2, 1}.
  assert(t.query_sum(0, 4) == 9);
  assert(t.query_max(0, 3) == 2);

  SegTreeBeats<int> lowest(1, numeric_limits<int>::lowest() + 1);
  lowest.chmin(0, 0, numeric_limits<int>::lowest());
  assert(lowest.at(0) == numeric_limits<int>::lowest());

  // Cross-check against a brute-force array over random clamps and queries.
  vector<int64_t> b{7, 1, 4, 9, 2, 6, 5, 8, 3, 0};
  SegTreeBeats<int64_t> s(b.begin(), b.end());
  int clamps[][3] = {{0, 9, 6}, {2, 5, 3}, {4, 8, 7}, {0, 4, 2}, {3, 9, 5}};
  for (auto &c : clamps) {
    for (int i = c[0]; i <= c[1]; i++) {
      b[i] = min(b[i], static_cast<int64_t>(c[2]));
    }
    s.chmin(c[0], c[1], c[2]);
    for (int lo = 0; lo < 10; lo++) {
      int64_t want_sum = 0, want_max = INT64_MIN;
      for (int hi = lo; hi < 10; hi++) {
        want_sum += b[hi];
        want_max = max(want_max, b[hi]);
        assert(s.query_sum(lo, hi) == want_sum);
        assert(s.query_max(lo, hi) == want_max);
      }
    }
  }
  return 0;
}