Use this to learn the idea, then write your own version.
1package main2 3import (4 "bufio"5 . "fmt"6 "io"7 "math/bits"8)9 1011type data01 struct{ cnt, sumC, sumC2 int }12type seg01 []struct {13 l, r int14 data0115 todo int16}17 18func (seg01) merge(l, r data01) data01 {19 return data01{l.cnt + r.cnt, l.sumC + r.sumC, l.sumC2 + r.sumC2}20}21 22func (t seg01) apply(o, f int) {23 cur := &t[o]24 cur.sumC2 += cur.sumC*f*2 + cur.cnt*f*f25 cur.sumC += cur.cnt * f26 cur.todo += f27}28 29func (t seg01) maintain(o int) {30 t[o].data01 = t.merge(t[o<<1].data01, t[o<<1|1].data01)31}32 33func (t seg01) spread(o int) {34 f := t[o].todo35 if f == 0 {36 return37 }38 t.apply(o<<1, f)39 t.apply(o<<1|1, f)40 t[o].todo = 041}42 43func (t seg01) build(o, l, r int) {44 t[o].l, t[o].r = l, r45 if l == r {46 return47 }48 m := (l + r) >> 149 t.build(o<<1, l, m)50 t.build(o<<1|1, m+1, r)51}52 53func (t seg01) set(o, i, cnt int) {54 if t[o].l == t[o].r {55 if cnt < 0 {56 t[o].data01 = data01{}57 } else {58 t[o].data01 = data01{1, cnt, cnt * cnt}59 }60 return61 }62 t.spread(o)63 m := (t[o].l + t[o].r) >> 164 if i <= m {65 t.set(o<<1, i, cnt)66 } else {67 t.set(o<<1|1, i, cnt)68 }69 t.maintain(o)70}71 72func (t seg01) update(o, l, r, f int) {73 if l <= t[o].l && t[o].r <= r {74 t.apply(o, f)75 return76 }77 t.spread(o)78 m := (t[o].l + t[o].r) >> 179 if l <= m {80 t.update(o<<1, l, r, f)81 }82 if m < r {83 t.update(o<<1|1, l, r, f)84 }85 t.maintain(o)86}87 88func (t seg01) query(o, l, r int) data01 {89 if l <= t[o].l && t[o].r <= r {90 return t[o].data0191 }92 t.spread(o)93 m := (t[o].l + t[o].r) >> 194 if r <= m {95 return t.query(o<<1, l, r)96 }97 if l > m {98 return t.query(o<<1|1, l, r)99 }100 return t.merge(t.query(o<<1, l, r), t.query(o<<1|1, l, r))101}102 103func cf1701F(in io.Reader, _w io.Writer) {104 out := bufio.NewWriter(_w)105 defer out.Flush()106 107 const mx int = 2e5108 t := make(seg01, 2<<bits.Len(uint(mx)))109 t.build(1, 1, mx)110 has := [mx + 1]bool{}111 112 var q, d, i int113 Fscan(in, &q, &d)114 for range q {115 Fscan(in, &i)116 if !has[i] {117 has[i] = true118 cnt := t.query(1, i, min(i+d, mx)).cnt119 t.set(1, i, cnt)120 if i > 1 {121 t.update(1, max(i-d, 1), i-1, 1)122 }123 } else {124 has[i] = false125 t.set(1, i, -1)126 if i > 1 {127 t.update(1, max(i-d, 1), i-1, -1)128 }129 }130 Fprintln(out, (t[1].sumC2-t[1].sumC)/2)131 }132}133 134135