///|
/// A monoid segment tree. op and e must be pure; op need not commute.
pub struct SegTree[S] {
priv n : Int
priv size : Int
priv data : Array[S]
priv op : (S, S) -> S
priv e : () -> S
}
///|
pub impl[S : Debug] Debug for SegTree[S] with fn to_repr(self) {
@debug.Repr((self.n, self.data))
}
///|
pub fn[S] SegTree::new(n : Int, op~ : (S, S) -> S, e~ : () -> S) -> SegTree[S] {
let (size, _) = @internal.tree_size(n)
{ n, size, data: Array::makei(2 * size, _ => e()), op, e, }
}
///|
pub fn[S] SegTree::from_array(
values : Array[S],
op~ : (S, S) -> S,
e~ : () -> S,
) -> SegTree[S] {
let tree = SegTree::new(values.length(), op~, e~)
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] SegTree::update(self : SegTree[S], k : Int) -> Unit {
self.data[k] = (self.op)(self.data[2 * k], self.data[2 * k + 1])
}
///|
pub fn[S] SegTree::set(self : SegTree[S], p : Int, value : S) -> Unit {
@internal.require(0 <= p && p < self.n)
let mut k = p + self.size
self.data[k] = value
while k > 1 {
k = k >> 1
self.update(k)
}
}
///|
pub fn[S] SegTree::get(self : SegTree[S], p : Int) -> S {
@internal.require(0 <= p && p < self.n)
self.data[self.size + p]
}
///|
pub fn[S] SegTree::prod(self : SegTree[S], l : Int, r : Int) -> S {
@internal.require(0 <= l && l <= r && r <= self.n)
let mut l = l + self.size
let mut r = r + self.size
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] SegTree::all_prod(self : SegTree[S]) -> S {
self.data[1]
}
///|
/// Requires predicate(e()) and a pure monotone predicate. O(log n).
pub fn[S] SegTree::max_right(
self : SegTree[S],
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
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 {
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
}
///|
/// Requires predicate(e()) and a pure monotone predicate. O(log n).
pub fn[S] SegTree::min_left(
self : SegTree[S],
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
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 {
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
}