Segment Tree

Share:
  1type S struct{ v int }
  2
  3func op(a, b S) S {
  4	if a.v > b.v {
  5		return a
  6	}
  7	return b
  8}
  9
 10func e() S {
 11	return S{0}
 12}
 13
 14type segTree struct {
 15	d  []S
 16	n  int
 17	sz int
 18	e  func() S
 19	f  func(a, b S) S
 20}
 21
 22func newSegTree(n int, e func() S, op func(a, b S) S) *segTree {
 23	sz := 1
 24	for sz < n {
 25		sz <<= 1
 26	}
 27	s := segTree{make([]S, 2*sz), n, sz, e, op}
 28	for i := 0; i < 2*sz; i++ {
 29		s.d[i] = e()
 30	}
 31	return &s
 32}
 33
 34func (s *segTree) Build(arr []S) {
 35	for i, v := range arr {
 36		s.d[i+s.sz] = v
 37	}
 38	for i := s.sz - 1; i > 0; i-- {
 39		s.d[i] = s.f(s.d[2*i], s.d[2*i+1])
 40	}
 41}
 42
 43func (s *segTree) Get(i int) S {
 44	return s.d[i+s.sz]
 45}
 46
 47func (s *segTree) Set(i int, a S) {
 48	i += s.sz
 49	s.d[i] = a
 50	for i >>= 1; i > 0; i >>= 1 {
 51		s.d[i] = s.f(s.d[2*i], s.d[2*i+1])
 52	}
 53}
 54
 55func (s *segTree) Query(l, r int) S {
 56	vl, vr := s.e(), s.e()
 57	l += s.sz
 58	r += s.sz
 59	for l < r {
 60		if l&1 == 1 {
 61			vl = s.f(vl, s.d[l])
 62			l++
 63		}
 64		if r&1 == 1 {
 65			r--
 66			vr = s.f(s.d[r], vr)
 67		}
 68		l >>= 1
 69		r >>= 1
 70	}
 71	return s.f(vl, vr)
 72}
 73
 74func (s *segTree) MaxRight(l int, f func(v S) bool) int {
 75	if l == s.n {
 76		return s.n
 77	}
 78	v := s.e()
 79	l += s.sz
 80
 81	for {
 82		for ; l&1 == 0; l >>= 1 {
 83		}
 84		if !f(s.f(v, s.d[l])) {
 85			for l < s.sz {
 86				l <<= 1
 87				if f(s.f(v, s.d[l])) {
 88					v = s.f(v, s.d[l])
 89					l++
 90				}
 91			}
 92			return l - s.sz
 93		}
 94		v = s.f(v, s.d[l])
 95		l++
 96		if l&-l == l {
 97			break
 98		}
 99	}
100
101	return s.n
102}
103
104func (s *segTree) MinLeft(r int, f func(v S) bool) int {
105	if r == 0 {
106		return 0
107	}
108	v := s.e()
109	r += s.sz
110
111	for {
112		for r--; r > 1 && r&1 == 1; r >>= 1 {
113		}
114		if !f(s.f(v, s.d[r])) {
115			for r < s.sz {
116				r = r<<1 + 1
117				if f(s.f(v, s.d[r])) {
118					v = s.f(v, s.d[r])
119					r--
120				}
121			}
122			return r + 1 - s.sz
123		}
124		v = s.f(v, s.d[r])
125		if r&-r == r {
126			break
127		}
128	}
129
130	return 0
131}
132