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) } }