- Translate each rule into one explicit state update.
- Maintain the invariant after every processed item.
- Return the accumulated state once all relevant input has been handled.
Code notes
- 186 lines of Go from the credited upstream file 1841E.go.
- The implementation visibly relies on sequence storage.
- No explicit loop blocks detected.
Complexity
Count the number and nesting of passes over the input, then include the maintained containers in the memory estimate.
Check the problem constraints before deciding whether this complexity will pass.
Use this to learn the idea, then write your own version.
1package main2 3import (4 "bufio"5 . "fmt"6 "io"7 "time"8)9 1011type node41 struct {12 lr [2]*node4113 priority uint14 key int15 subCnt int16}17 18func (o *node41) size() int {19 if o != nil {20 return o.subCnt 21 }22 return 023}24 25func (o *node41) maintain() { o.subCnt = 1 + o.lr[0].size() + o.lr[1].size() }26 27func (o *node41) rotate(d int) *node41 {28 x := o.lr[d^1]29 o.lr[d^1] = x.lr[d]30 x.lr[d] = o31 o.maintain()32 x.maintain()33 return x34}35 36type treap41 struct {37 rd uint38 root *node4139}40 41func (t *treap41) fastRand() uint {42 t.rd ^= t.rd << 1343 t.rd ^= t.rd >> 1744 t.rd ^= t.rd << 545 return t.rd46}47 48func (t *treap41) size() int { return t.root.size() }49 50func (t *treap41) _put(o *node41, key int) *node41 {51 if o == nil {52 return &node41{priority: t.fastRand(), key: key, subCnt: 1}53 }54 if d := o.cmp(key); d >= 0 {55 o.lr[d] = t._put(o.lr[d], key)56 if o.lr[d].priority > o.priority {57 o = o.rotate(d ^ 1)58 }59 } else {60 61 }62 o.maintain()63 return o64}65 66func (t *treap41) put(key int) { t.root = t._put(t.root, key) }67 68func (t *treap41) _delete(o *node41, key int) *node41 {69 if o == nil {70 return nil71 }72 if d := o.cmp(key); d >= 0 {73 o.lr[d] = t._delete(o.lr[d], key)74 } else {75 if o.lr[1] == nil {76 return o.lr[0]77 }78 if o.lr[0] == nil {79 return o.lr[1]80 }81 d = 082 if o.lr[0].priority > o.lr[1].priority {83 d = 184 }85 o = o.rotate(d)86 o.lr[d] = t._delete(o.lr[d], key)87 }88 o.maintain()89 return o90}91 92func (t *treap41) delete(key int) { t.root = t._delete(t.root, key) }93 94func (o *node41) cmp(a int) int {95 b := o.key96 if a == b {97 return -198 }99 if a < b {100 return 0101 }102 return 1103}104 105func (t *treap41) lowerBound(key int) (lb *node41) {106 for o := t.root; o != nil; {107 switch c := o.cmp(key); {108 case c == 0:109 lb = o110 o = o.lr[0]111 case c > 0:112 o = o.lr[1]113 default:114 return o115 }116 }117 return118}119 120func (t *treap41) prev(key int) (prev *node41) {121 for o := t.root; o != nil; {122 if o.cmp(key) <= 0 {123 o = o.lr[0]124 } else {125 prev = o126 o = o.lr[1]127 }128 }129 return130}131 132func CF1841E(_r io.Reader, _w io.Writer) {133 in := bufio.NewReader(_r)134 out := bufio.NewWriter(_w)135 defer out.Flush()136 min := func(a, b int64) int64 {137 if a > b {138 return b139 }140 return a141 }142 143 var T, n, v int144 var m int64145 for Fscan(in, &T); T > 0; T-- {146 Fscan(in, &n)147 col := make([][]int, n+1)148 for i := 0; i < n; i++ {149 Fscan(in, &v)150 col[v] = append(col[v], i)151 }152 153 cnt := make([]int, n+1)154 cnt[n] = n155 t := &treap41{rd: uint(time.Now().UnixNano())/2 + 1}156 t.put(-1)157 t.put(n)158 for i := n; i > 0; i-- {159 for _, j := range col[i] {160 r := t.lowerBound(j).key161 l := t.prev(r).key162 cnt[r-l-1] -= i163 cnt[j-l-1] += i164 cnt[r-j-1] += i165 t.put(j)166 }167 }168 169 ans := int64(0)170 Fscan(in, &m)171 for i := int64(n); i > 1 && m > 1; i-- {172 c := min(m/i, int64(cnt[i]))173 ans += c * (i - 1)174 m -= c * i175 cnt[i] -= int(c)176 if 1 < m && m < i && cnt[i] > 0 {177 ans += m - 1178 break179 }180 }181 Fprintln(out, ans)182 }183}184 185186