Lazy Segment Tree

Share:
  • func newLazySegmentTree(n int, e func() S, id func() F, f func(a, b S) S, g func(a S, x F) S, h func(a, x F) F,) *lazySegmentTree
  • func (s *lazySegmentTree) Set(i int, a S)
  • func (s *lazySegmentTree) Build()
  • func (s *lazySegmentTree) Get(k int) S
  • func (s *lazySegmentTree) Update(l, r int, x F)
  • func (s *lazySegmentTree) Query(l, r int) S
  • func (s *lazySegmentTree) MaxRight(l int, f func(v S) bool) int
  • func (s *lazySegmentTree) MinLeft(r int, f func(v S) bool) int
  1type S struct{}
  2type F struct{}
  3
  4type lazySegmentTree struct {
  5	d  []S
  6	lz []F
  7	n  int
  8	sz int
  9	lg int
 10	e  func() S
 11	id func() F
 12	f  func(a, b S) S
 13	g  func(a S, x F) S
 14	h  func(a, x F) F
 15}
 16
 17func newLazySegmentTree(n int,
 18	e func() S,
 19	id func() F,
 20	f func(a, b S) S,
 21	g func(a S, x F) S,
 22	h func(a, x F) F,
 23) *lazySegmentTree {
 24	sz, lg := 1, 0
 25	for sz < n {
 26		sz <<= 1
 27		lg++
 28	}
 29	d := make([]S, 2*sz)
 30	lz := make([]F, 2*sz)
 31	for i := 0; i < 2*sz; i++ {
 32		d[i] = e()
 33		lz[i] = id()
 34	}
 35	s := lazySegmentTree{d, lz, n, sz, lg, e, id, f, g, h}
 36	return &s
 37}
 38
 39func (s *lazySegmentTree) Set(i int, a S) {
 40	s.d[i+s.sz] = a
 41}
 42
 43func (s *lazySegmentTree) Build() {
 44	for i := s.sz - 1; i > 0; i-- {
 45		s.d[i] = s.f(s.d[2*i], s.d[2*i+1])
 46	}
 47}
 48
 49func (s *lazySegmentTree) Get(k int) S {
 50	k += s.sz
 51	for i := s.lg; i > 0; i-- {
 52		s.prop(k >> i)
 53	}
 54	return s.d[k]
 55}
 56
 57func (s *lazySegmentTree) Update(l, r int, x F) {
 58	l += s.sz
 59	r += s.sz
 60	for i := s.lg; i > 0; i-- {
 61		if l>>i<<i != l {
 62			s.prop(l >> i)
 63		}
 64		if r>>i<<i != r {
 65			s.prop((r - 1) >> i)
 66		}
 67	}
 68
 69	for ll, rr := l, r; ll < rr; ll, rr = ll>>1, rr>>1 {
 70		if ll&1 == 1 {
 71			s.apply(ll, x)
 72			ll++
 73		}
 74		if rr&1 == 1 {
 75			rr--
 76			s.apply(rr, x)
 77		}
 78	}
 79
 80	for i := 1; i <= s.lg; i++ {
 81		if l>>i<<i != l {
 82			k := l >> i
 83			s.d[k] = s.f(s.d[2*k], s.d[2*k+1])
 84		}
 85		if r>>i<<i != r {
 86			k := (r - 1) >> i
 87			s.d[k] = s.f(s.d[2*k], s.d[2*k+1])
 88		}
 89	}
 90}
 91
 92func (s *lazySegmentTree) Query(l, r int) S {
 93	vl, vr := s.e(), s.e()
 94	l += s.sz
 95	r += s.sz
 96
 97	for i := s.lg; i > 0; i-- {
 98		if l>>i<<i != l {
 99			s.prop(l >> i)
100		}
101		if r>>i<<i != r {
102			s.prop(r >> i)
103		}
104	}
105
106	for ; l < r; l, r = l>>1, r>>1 {
107		if l&1 == 1 {
108			vl = s.f(vl, s.d[l])
109			l++
110		}
111		if r&1 == 1 {
112			r--
113			vr = s.f(s.d[r], vr)
114		}
115	}
116	return s.f(vl, vr)
117}
118
119func (s *lazySegmentTree) MaxRight(l int, f func(v S) bool) int {
120	if l == s.n {
121		return s.n
122	}
123	v := s.e()
124	l += s.sz
125	for i := s.lg; i > 0; i-- {
126		s.prop(l >> i)
127	}
128
129	for {
130		for l&1 == 0 {
131			l >>= 1
132		}
133		if !f(s.f(v, s.d[l])) {
134			for l < s.sz {
135				l <<= 1
136				if f(s.f(v, s.d[l])) {
137					v = s.f(v, s.d[l])
138					l++
139				}
140			}
141			return l - s.sz
142		}
143		v = s.f(v, s.d[l])
144		l++
145		if l&-l == l {
146			break
147		}
148	}
149
150	return s.n
151}
152
153func (s *lazySegmentTree) MinLeft(r int, f func(v S) bool) int {
154	if r == 0 {
155		return 0
156	}
157	v := s.e()
158	r += s.sz
159	for i := s.lg; i > 0; i-- {
160		s.prop((r - 1) >> i)
161	}
162
163	for {
164		r--
165		for r > 1 && r&1 == 1 {
166			r >>= 1
167		}
168		if !f(s.f(v, s.d[r])) {
169			for r < s.sz {
170				r = r<<1 + 1
171				if f(s.f(v, s.d[r])) {
172					v = s.f(v, s.d[r])
173					r--
174				}
175			}
176			return r + 1 - s.sz
177		}
178		v = s.f(v, s.d[r])
179		if r&-r == r {
180			break
181		}
182	}
183
184	return 0
185}
186
187func (s *lazySegmentTree) prop(k int) {
188	s.apply(2*k, s.lz[k])
189	s.apply(2*k+1, s.lz[k])
190	s.lz[k] = s.id()
191}
192
193func (s *lazySegmentTree) apply(k int, x F) {
194	s.d[k] = s.g(s.d[k], x)
195	if k < s.sz {
196		s.lz[k] = s.h(s.lz[k], x)
197	}
198}
199