///|
/// Lazy monoid tree. composition(f,g) means f after g. All callbacks must be pure.
pub struct LazySegTree[S, F] {
priv n : Int
priv size : Int
priv log : Int
priv data : Array[S]
priv pending : Array[F]
priv op : (S, S) -> S
priv e : () -> S
priv mapping : (F, S) -> S
priv composition : (F, F) -> F
priv id : () -> F
}
///|
pub impl[S : Debug, F : Debug] Debug for LazySegTree[S, F] with fn to_repr(self) {
@debug.Repr((self.n, self.data, self.pending))
}
///|
pub fn[S, F] LazySegTree::new(
n : Int,
op~ : (S, S) -> S,
e~ : () -> S,
mapping~ : (F, S) -> S,
composition~ : (F, F) -> F,
id~ : () -> F,
) -> LazySegTree[S, F] {
let (size, log) = @internal.tree_size(n)
{
n,
size,
log,
data: Array::makei(2 * size, _ => e()),
pending: Array::makei(size, _ => id()),
op,
e,
mapping,
composition,
id,
}
}
///|
pub fn[S, F] LazySegTree::from_array(
values : Array[S],
op~ : (S, S) -> S,
e~ : () -> S,
mapping~ : (F, S) -> S,
composition~ : (F, F) -> F,
id~ : () -> F,
) -> LazySegTree[S, F] {
let tree = LazySegTree::new(
values.length(),
op~,
e~,
mapping~,
composition~,
id~,
)
for i = 0; i < values.length(); i = i + 1 {
tree.data[tree.size + i] = values[i]
}
for i = tree.size - 1; i > 0; i = i - 1 {
tree.update(i)
}
tree
}
///|
fn[S, F] LazySegTree::update(self : LazySegTree[S, F], k : Int) -> Unit {
self.data[k] = (self.op)(self.data[2 * k], self.data[2 * k + 1])
}
///|
fn[S, F] LazySegTree::all_apply(
self : LazySegTree[S, F],
k : Int,
f : F,
) -> Unit {
self.data[k] = (self.mapping)(f, self.data[k])
if k < self.size {
self.pending[k] = (self.composition)(f, self.pending[k])
}
}
///|
fn[S, F] LazySegTree::push(self : LazySegTree[S, F], k : Int) -> Unit {
self.all_apply(2 * k, self.pending[k])
self.all_apply(2 * k + 1, self.pending[k])
self.pending[k] = (self.id)()
}
///|
pub fn[S, F] LazySegTree::set(
self : LazySegTree[S, F],
p : Int,
value : S,
) -> Unit {
@internal.require(0 <= p && p < self.n)
let p = p + self.size
for i = self.log; i >= 1; i = i - 1 {
self.push(p >> i)
}
self.data[p] = value
for i = 1; i <= self.log; i = i + 1 {
self.update(p >> i)
}
}
///|
pub fn[S, F] LazySegTree::get(self : LazySegTree[S, F], p : Int) -> S {
@internal.require(0 <= p && p < self.n)
let p = p + self.size
for i = self.log; i >= 1; i = i - 1 {
self.push(p >> i)
}
self.data[p]
}
///|
pub fn[S, F] LazySegTree::prod(self : LazySegTree[S, F], l : Int, r : Int) -> S {
@internal.require(0 <= l && l <= r && r <= self.n)
if l == r {
return (self.e)()
}
let mut l = l + self.size
let mut r = r + self.size
for i = self.log; i >= 1; i = i - 1 {
if l >> i << i != l {
self.push(l >> i)
}
if r >> i << i != r {
self.push((r - 1) >> i)
}
}
let mut left = (self.e)()
let mut right = (self.e)()
while l < r {
if (l & 1) != 0 {
left = (self.op)(left, self.data[l])
l = l + 1
}
if (r & 1) != 0 {
r -= 1
right = (self.op)(self.data[r], right)
}
l = l >> 1
r = r >> 1
}
(self.op)(left, right)
}
///|
pub fn[S, F] LazySegTree::all_prod(self : LazySegTree[S, F]) -> S {
self.data[1]
}
///|
pub fn[S, F] LazySegTree::apply(
self : LazySegTree[S, F],
p : Int,
f : F,
) -> Unit {
@internal.require(0 <= p && p < self.n)
let p = p + self.size
for i = self.log; i >= 1; i = i - 1 {
self.push(p >> i)
}
self.data[p] = (self.mapping)(f, self.data[p])
for i = 1; i <= self.log; i = i + 1 {
self.update(p >> i)
}
}
///|
pub fn[S, F] LazySegTree::apply_range(
self : LazySegTree[S, F],
l : Int,
r : Int,
f : F,
) -> Unit {
@internal.require(0 <= l && l <= r && r <= self.n)
if l == r {
return
}
let l0 = l + self.size
let r0 = r + self.size
for i = self.log; i >= 1; i = i - 1 {
if l0 >> i << i != l0 {
self.push(l0 >> i)
}
if r0 >> i << i != r0 {
self.push((r0 - 1) >> i)
}
}
let mut l = l0
let mut r = r0
while l < r {
if (l & 1) != 0 {
self.all_apply(l, f)
l = l + 1
}
if (r & 1) != 0 {
r -= 1
self.all_apply(r, f)
}
l = l >> 1
r = r >> 1
}
for i = 1; i <= self.log; i = i + 1 {
if l0 >> i << i != l0 {
self.update(l0 >> i)
}
if r0 >> i << i != r0 {
self.update((r0 - 1) >> i)
}
}
}
///|
pub fn[S, F] LazySegTree::max_right(
self : LazySegTree[S, F],
l : Int,
predicate : (S) -> Bool,
) -> Int {
@internal.require(0 <= l && l <= self.n && predicate((self.e)()))
if l == self.n {
return self.n
}
let mut l = l + self.size
for i = self.log; i >= 1; i = i - 1 {
self.push(l >> i)
}
let mut acc = (self.e)()
while true {
while (l & 1) == 0 {
l = l >> 1
}
if !predicate((self.op)(acc, self.data[l])) {
while l < self.size {
self.push(l)
l = l * 2
let next = (self.op)(acc, self.data[l])
if predicate(next) {
acc = next
l = l + 1
}
}
return l - self.size
}
acc = (self.op)(acc, self.data[l])
l = l + 1
if (l & -l) == l {
break
}
}
self.n
}
///|
pub fn[S, F] LazySegTree::min_left(
self : LazySegTree[S, F],
r : Int,
predicate : (S) -> Bool,
) -> Int {
@internal.require(0 <= r && r <= self.n && predicate((self.e)()))
if r == 0 {
return 0
}
let mut r = r + self.size
for i = self.log; i >= 1; i = i - 1 {
self.push((r - 1) >> i)
}
let mut acc = (self.e)()
while true {
r -= 1
while r > 1 && (r & 1) != 0 {
r = r >> 1
}
if !predicate((self.op)(self.data[r], acc)) {
while r < self.size {
self.push(r)
r = r * 2 + 1
let next = (self.op)(self.data[r], acc)
if predicate(next) {
acc = next
r = r - 1
}
}
return r + 1 - self.size
}
acc = (self.op)(self.data[r], acc)
if (r & -r) == r {
break
}
}
0
}