package main import ( "database/sql" ) type segmentTree struct { N int seg []int } // Build segment tree func (s *segmentTree) Build(a *sql.Rows, l, r, n int) { if l == r { a.Next() a.Scan(&s.seg[n]) return } m := (l + r) >> 1 s.Build(a, l, m, n<<1) s.Build(a, m+1, r, n<<1|1) s.seg[n] = s.seg[n<<1] + s.seg[n<<1|1] } // Update value at index x func (s *segmentTree) Update(x, v, l, r, n int) { if l == r { s.seg[n] = v return } m := (l + r) >> 1 if x <= m { s.Update(x, v, l, m, n<<1) } else { s.Update(x, v, m+1, r, n<<1|1) } s.seg[n] = s.seg[n<<1] + s.seg[n<<1|1] } // Find element with prefix sum v func (s *segmentTree) Query(v, l, r, n int) (int, int) { if l == r { return s.seg[n], l } m := (l + r) >> 1 if s.seg[n<<1] >= v { return s.Query(v, l, m, n<<1) } else { return s.Query(v-s.seg[n<<1], m+1, r, n<<1|1) } }