Use this to learn the idea, then write your own version.
1package main2 3import (4 "bufio"5 "cmp"6 . "fmt"7 "io"8 "math/bits"9 "time"10)11 1213type nodeM76[K comparable, V any] struct {14 son [2]*nodeM76[K, V]15 priority uint16 key K17 value V18 subSize int19}20 21func (o *nodeM76[K, V]) size() int {22 if o != nil {23 return o.subSize24 }25 return 026}27 28func (o *nodeM76[K, V]) maintain() {29 o.subSize = 1 + o.son[0].size() + o.son[1].size()30}31 32func (o *nodeM76[K, V]) rotate(d int) *nodeM76[K, V] {33 x := o.son[d^1]34 o.son[d^1] = x.son[d]35 x.son[d] = o36 o.maintain()37 x.maintain()38 return x39}40 41type treapM76[K comparable, V any] struct {42 rd uint43 root *nodeM76[K, V]44 comparator func(a, b K) int45}46 47func (t *treapM76[K, V]) fastRand() uint {48 t.rd ^= t.rd << 1349 t.rd ^= t.rd >> 1750 t.rd ^= t.rd << 551 return t.rd52}53 54func (t *treapM76[K, V]) size() int { return t.root.size() }55func (t *treapM76[K, V]) empty() bool { return t.size() == 0 }56 57func (t *treapM76[K, V]) _put(o *nodeM76[K, V], key K, value V) *nodeM76[K, V] {58 if o == nil {59 o = &nodeM76[K, V]{priority: t.fastRand(), key: key, value: value}60 } else {61 c := t.comparator(key, o.key)62 if c == 0 {63 o.value = value64 } else {65 d := 066 if c > 0 {67 d = 168 }69 o.son[d] = t._put(o.son[d], key, value)70 if o.son[d].priority > o.priority {71 o = o.rotate(d ^ 1)72 }73 }74 }75 o.maintain()76 return o77}78 79func (t *treapM76[K, V]) put(key K, value V) { t.root = t._put(t.root, key, value) }80 81func (t *treapM76[K, V]) _delete(o *nodeM76[K, V], key K) *nodeM76[K, V] {82 if o == nil {83 return nil84 }85 if c := t.comparator(key, o.key); c != 0 {86 d := 087 if c > 0 {88 d = 189 }90 o.son[d] = t._delete(o.son[d], key)91 } else {92 if o.son[1] == nil {93 return o.son[0]94 }95 if o.son[0] == nil {96 return o.son[1]97 }98 d := 099 if o.son[0].priority > o.son[1].priority {100 d = 1101 }102 o = o.rotate(d)103 o.son[d] = t._delete(o.son[d], key)104 }105 o.maintain()106 return o107}108 109func (t *treapM76[K, V]) delete(key K) { t.root = t._delete(t.root, key) }110 111func (t *treapM76[K, V]) min() *nodeM76[K, V] { return t.kth(0) }112func (t *treapM76[K, V]) max() *nodeM76[K, V] { return t.kth(t.size() - 1) }113 114func (t *treapM76[K, V]) lowerBoundIndex(key K) (kth int) {115 for o := t.root; o != nil; {116 c := t.comparator(key, o.key)117 if c < 0 {118 o = o.son[0]119 } else if c > 0 {120 kth += o.son[0].size() + 1121 o = o.son[1]122 } else {123 kth += o.son[0].size()124 break125 }126 }127 return128}129 130func (t *treapM76[K, V]) upperBoundIndex(key K) (kth int) {131 for o := t.root; o != nil; {132 c := t.comparator(key, o.key)133 if c < 0 {134 o = o.son[0]135 } else if c > 0 {136 kth += o.son[0].size() + 1137 o = o.son[1]138 } else {139 kth += o.son[0].size() + 1140 break141 }142 }143 return144}145 146func (t *treapM76[K, V]) kth(k int) (o *nodeM76[K, V]) {147 if k < 0 || k >= t.root.size() {148 return149 }150 for o = t.root; o != nil; {151 leftSize := o.son[0].size()152 if k < leftSize {153 o = o.son[0]154 } else {155 k -= leftSize + 1156 if k < 0 {157 break158 }159 o = o.son[1]160 }161 }162 return163}164 165func (t *treapM76[K, V]) prev(key K) *nodeM76[K, V] { return t.kth(t.lowerBoundIndex(key) - 1) }166func (t *treapM76[K, V]) next(key K) *nodeM76[K, V] { return t.kth(t.upperBoundIndex(key)) }167 168func newMap76[K cmp.Ordered, V any]() *treapM76[K, V] {169 return &treapM76[K, V]{170 rd: uint(time.Now().UnixNano()),171 comparator: cmp.Compare[K],172 }173}174 175func cf176E(in io.Reader, _w io.Writer) {176 out := bufio.NewWriter(_w)177 defer out.Flush()178 var n, q, v, w, wt, ans, ts int179 var op string180 Fscan(in, &n)181 type nb struct{ to, wt int }182 g := make([][]nb, n)183 for range n - 1 {184 Fscan(in, &v, &w, &wt)185 v--186 w--187 g[v] = append(g[v], nb{w, wt})188 g[w] = append(g[w], nb{v, wt})189 }190 191 const mx = 17192 pa := make([][mx]int, n)193 dep := make([]int, n)194 dis := make([]int, n)195 dfn := make([]int, n)196 var dfs func(int, int)197 dfs = func(v, p int) {198 ts++199 dfn[v] = ts200 pa[v][0] = p201 for _, e := range g[v] {202 w := e.to203 if w == p {204 continue205 }206 dep[w] = dep[v] + 1207 dis[w] = dis[v] + e.wt208 dfs(w, v)209 }210 }211 dfs(0, -1)212 for i := range mx - 1 {213 for v := range pa {214 if p := pa[v][i]; p != -1 {215 pa[v][i+1] = pa[p][i]216 } else {217 pa[v][i+1] = -1218 }219 }220 }221 uptoDep := func(v, d int) int {222 for k := uint(dep[v] - d); k > 0; k &= k - 1 {223 v = pa[v][bits.TrailingZeros(k)]224 }225 return v226 }227 getLCA := func(v, w int) int {228 if dep[v] > dep[w] {229 v, w = w, v230 }231 w = uptoDep(w, dep[v])232 if w == v {233 return v234 }235 for i := mx - 1; i >= 0; i-- {236 if pv, pw := pa[v][i], pa[w][i]; pv != pw {237 v, w = pv, pw238 }239 }240 return pa[v][0]241 }242 getDis := func(v, w int) int { return dis[v] + dis[w] - dis[getLCA(v, w)]*2 }243 244 t := newMap76[int, int]()245 do := func(d, mul int) {246 if t.empty() {247 return248 }249 o := t.prev(d)250 if o == nil {251 o = t.max()252 }253 pre := o.value254 o = t.next(d)255 if o == nil {256 o = t.min()257 }258 nxt := o.value259 ans += (getDis(pre, v) + getDis(v, nxt) - getDis(pre, nxt)) * mul260 }261 262 Fscan(in, &q)263 for range q {264 Fscan(in, &op)265 if op == "?" {266 Fprintln(out, ans/2)267 continue268 }269 Fscan(in, &v)270 v--271 d := dfn[v]272 if op == "+" {273 do(d, 1)274 t.put(d, v)275 } else {276 t.delete(d)277 do(d, -1)278 }279 }280}281 282283