Skip to content

Commit 452f9f1

Browse files
Added segtree
1 parent 0e8242f commit 452f9f1

File tree

2 files changed

+163
-2
lines changed

2 files changed

+163
-2
lines changed

src/lib.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ pub use dsu::Dsu;
2222
pub use fenwicktree::FenwickTree;
2323
pub use math::{crt, floor_sum, inv_mod, pow_mod};
2424
pub use mincostflow::MinCostFlowGraph;
25+
pub use segtree::{Monoid, Segtree};
2526
pub use string::{
2627
lcp_array, lcp_array_arbitrary, suffix_array, suffix_array_arbitrary, suffix_array_manual,
2728
z_algorithm, z_algorithm_arbitrary,

src/segtree.rs

Lines changed: 162 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,162 @@
1-
struct Segtree {
2-
}
1+
use crate::internal_bit::ceil_pow2;
2+
3+
pub trait Monoid {
4+
type S: Copy;
5+
const IDENTITY: Self::S;
6+
fn binary_operation(a: Self::S, b: Self::S) -> Self::S;
7+
}
8+
9+
impl<M: Monoid> Segtree<M> {
10+
pub fn new(n: usize) -> Segtree<M> {
11+
vec![M::IDENTITY; n].into()
12+
}
13+
}
14+
impl<M: Monoid> From<Vec<M::S>> for Segtree<M> {
15+
fn from(v: Vec<M::S>) -> Self {
16+
let n = v.len();
17+
let log = ceil_pow2(n as u32) as usize;
18+
let size = 1 << log;
19+
let mut d = vec![M::IDENTITY; n];
20+
d[size..(n + size)].clone_from_slice(&v[..n]);
21+
let mut ret = Segtree { n, size, log, d };
22+
for i in (1..n).rev() {
23+
ret.update(i);
24+
}
25+
ret
26+
}
27+
}
28+
29+
impl<M: Monoid> Segtree<M> {
30+
pub fn set(&mut self, mut p: usize, x: M::S) {
31+
assert!(p <= self.n);
32+
p += self.size;
33+
self.d[p] = x;
34+
for i in 1..=self.log {
35+
self.update(p >> i);
36+
}
37+
}
38+
39+
pub fn get(&self, p: usize) -> M::S {
40+
assert!(p <= self.n);
41+
self.d[p + self.size]
42+
}
43+
44+
pub fn prod(&self, mut l: usize, mut r: usize) -> M::S {
45+
assert!(l <= r && r <= self.n);
46+
let mut sml = M::IDENTITY;
47+
let mut smr = M::IDENTITY;
48+
l += self.size;
49+
r += self.size;
50+
51+
while l < r {
52+
if l & 1 != 0 {
53+
sml = M::binary_operation(sml, self.d[l]);
54+
l += 1;
55+
}
56+
if r & 1 != 0 {
57+
r -= 1;
58+
smr = M::binary_operation(self.d[r], smr);
59+
}
60+
l >>= 1;
61+
r >>= 1;
62+
}
63+
64+
M::binary_operation(sml, smr)
65+
}
66+
67+
pub fn all_prod(&self) -> M::S {
68+
self.d[1]
69+
}
70+
71+
pub fn max_right<F>(&self, mut l: usize, f: F) -> usize
72+
where
73+
F: Fn(M::S) -> bool,
74+
{
75+
assert!(l <= self.n);
76+
assert!(f(M::IDENTITY));
77+
if l == self.n {
78+
return self.n;
79+
}
80+
l += self.size;
81+
let mut sm = M::IDENTITY;
82+
while {
83+
// do
84+
while l % 2 == 0 {
85+
l >>= 1;
86+
}
87+
if !f(M::binary_operation(sm, self.d[l])) {
88+
while l < self.size {
89+
l *= 2;
90+
let res = M::binary_operation(sm, self.d[l]);
91+
if f(res) {
92+
sm = res;
93+
l += 1;
94+
}
95+
}
96+
return l - self.size;
97+
}
98+
sm = M::binary_operation(sm, self.d[l]);
99+
l += 1;
100+
// while
101+
{
102+
let l = l as isize;
103+
(l & -l) != l
104+
}
105+
} {}
106+
self.n
107+
}
108+
109+
pub fn min_left<F>(&self, mut r: usize, f: F) -> usize
110+
where
111+
F: Fn(M::S) -> bool,
112+
{
113+
assert!(r <= self.n);
114+
assert!(f(M::IDENTITY));
115+
if r == 0 {
116+
return 0;
117+
}
118+
r += self.size;
119+
let mut sm = M::IDENTITY;
120+
while {
121+
// do
122+
r -= 1;
123+
while r > 1 && r % 2 == 1 {
124+
r >>= 1;
125+
}
126+
if !f(M::binary_operation(self.d[r], sm)) {
127+
while r < self.size {
128+
r = 2 * r + 1;
129+
let res = M::binary_operation(self.d[r], sm);
130+
if f(res) {
131+
sm = res;
132+
r -= 1;
133+
}
134+
}
135+
return r + 1 - self.size;
136+
}
137+
sm = M::binary_operation(self.d[r], sm);
138+
// while
139+
{
140+
let r = r as isize;
141+
(r & -r) != r
142+
}
143+
} {}
144+
0
145+
}
146+
147+
fn update(&mut self, k: usize) {
148+
self.d[k] = M::binary_operation(self.d[2 * k], self.d[2 * k + 1]);
149+
}
150+
}
151+
152+
#[derive(Default)]
153+
pub struct Segtree<M>
154+
where
155+
M: Monoid,
156+
{
157+
// variable name is _n in original library
158+
n: usize,
159+
size: usize,
160+
log: usize,
161+
d: Vec<M::S>,
162+
}

0 commit comments

Comments
 (0)