Codeforces 1398E — Two Types of Spells

1package main2 3import (4	"bufio"5	"container/heap"6	. "fmt"7	"io"8	"time"9)10 11// github.com/EndlessCheng/codeforces-go12type node98 struct {13	lr       [2]*node9814	priority uint15	key, sz  int16	s        int6417}18 19func (o *node98) cmp(b int) int {20	switch {21	case b < o.key:22		return 023	case b > o.key:24		return 125	default:26		return -127	}28}29 30func (o *node98) size() int {31	if o != nil {32		return o.sz33	}34	return 035}36 37func (o *node98) sum() int64 {38	if o != nil {39		return o.s40	}41	return 042}43 44func (o *node98) maintain() {45	o.sz = 1 + o.lr[0].size() + o.lr[1].size()46	o.s = int64(o.key) + o.lr[0].sum() + o.lr[1].sum()47}48 49func (o *node98) rotate(d int) *node98 {50	x := o.lr[d^1]51	o.lr[d^1] = x.lr[d]52	x.lr[d] = o53	o.maintain()54	x.maintain()55	return x56}57 58type treap98 struct {59	rd   uint60	root *node9861}62 63func (t *treap98) fastRand() uint {64	t.rd ^= t.rd << 1365	t.rd ^= t.rd >> 1766	t.rd ^= t.rd << 567	return t.rd68}69 70func (t *treap98) _put(o *node98, key int) *node98 {71	if o == nil {72		return &node98{priority: t.fastRand(), key: key, sz: 1, s: int64(key)}73	}74	if d := o.cmp(key); d >= 0 {75		o.lr[d] = t._put(o.lr[d], key)76		if o.lr[d].priority > o.priority {77			o = o.rotate(d ^ 1)78		}79	}80	o.maintain()81	return o82}83 84func (t *treap98) put(key int) { t.root = t._put(t.root, key) }85 86func (t *treap98) _delete(o *node98, key int) *node98 {87	if o == nil {88		return nil89	}90	if d := o.cmp(key); d >= 0 {91		o.lr[d] = t._delete(o.lr[d], key)92	} else {93		if o.lr[1] == nil {94			return o.lr[0]95		}96		if o.lr[0] == nil {97			return o.lr[1]98		}99		d = 0100		if o.lr[0].priority > o.lr[1].priority {101			d = 1102		}103		o = o.rotate(d)104		o.lr[d] = t._delete(o.lr[d], key)105	}106	o.maintain()107	return o108}109 110func (t *treap98) delete(key int) { t.root = t._delete(t.root, key) }111 112func (t *treap98) kth(k int) (s int64) {113	for o := t.root; o != nil; {114		if ls := o.lr[0].size(); k < ls {115			o = o.lr[0]116		} else {117			s += o.lr[0].sum()118			if k > ls {119				s += int64(o.key)120			}121			k -= ls + 1122			if k < 0 {123				return124			}125			o = o.lr[1]126		}127	}128	return129}130 131type vi98 struct{ v, i int }132type mh98 []*vi98133 134func (h mh98) Len() int            { return len(h) }135func (h mh98) Less(i, j int) bool  { return h[i].v < h[j].v }136func (h mh98) Swap(i, j int)       { h[i], h[j] = h[j], h[i]; h[i].i = i; h[j].i = j }137func (h *mh98) Push(v interface{}) { *h = append(*h, v.(*vi98)) }138func (h *mh98) Pop() interface{}   { a := *h; v := a[len(a)-1]; *h = a[:len(a)-1]; return v }139func (h *mh98) push(v int) *vi98   { p := &vi98{v, len(*h)}; heap.Push(h, p); return p }140 141func CF1398E(_r io.Reader, _w io.Writer) {142	in := bufio.NewReader(_r)143	out := bufio.NewWriter(_w)144	defer out.Flush()145 146	t := &treap98{rd: uint(time.Now().UnixNano())/2 + 1}147	h := mh98{}148	ptr := map[int]*vi98{}149 150	var q, tp, v int151	for Fscan(in, &q); q > 0; q-- {152		Fscan(in, &tp, &v)153		if v > 0 {154			t.put(v)155		} else {156			t.delete(-v)157		}158		if tp == 1 {159			if v > 0 {160				ptr[v] = h.push(v)161			} else {162				heap.Remove(&h, ptr[-v].i)163				delete(ptr, -v)164			}165		}166		if len(h) == 0 {167			Fprintln(out, t.root.sum())168		} else if len(h) == t.root.sz {169			Fprintln(out, t.root.sum()*2-int64(h[0].v))170		} else {171			mi := h[0].v172			t.delete(mi)173			extra := t.kth(t.root.sz - len(h))174			t.put(mi)175			Fprintln(out, t.root.sum()*2-extra-int64(mi))176		}177	}178}179 180//func main() { CF1398E(os.Stdin, os.Stdout) }181