Golang | Leetcode Golang题解之第327题区间和的个数

题目：

题解：

import "math/rand" // 默认导入的 rand 不是这个库，需要显式指明type node struct {ch       [2]*nodepriority intkey      intdupCnt   intsz       int
}func (o *node) cmp(b int) int {switch {case b < o.key:return 0case b > o.key:return 1default:return -1}
}func (o *node) size() int {if o != nil {return o.sz}return 0
}func (o *node) maintain() {o.sz = o.dupCnt + o.ch[0].size() + o.ch[1].size()
}func (o *node) rotate(d int) *node {x := o.ch[d^1]o.ch[d^1] = x.ch[d]x.ch[d] = oo.maintain()x.maintain()return x
}type treap struct {root *node
}func (t *treap) _insert(o *node, key int) *node {if o == nil {return &node{priority: rand.Int(), key: key, dupCnt: 1, sz: 1}}if d := o.cmp(key); d >= 0 {o.ch[d] = t._insert(o.ch[d], key)if o.ch[d].priority > o.priority {o = o.rotate(d ^ 1)}} else {o.dupCnt++}o.maintain()return o
}func (t *treap) insert(key int) {t.root = t._insert(t.root, key)
}// equal=false: 小于 key 的元素个数
// equal=true: 小于或等于 key 的元素个数
func (t *treap) rank(key int, equal bool) (cnt int) {for o := t.root; o != nil; {switch c := o.cmp(key); {case c == 0:o = o.ch[0]case c > 0:cnt += o.dupCnt + o.ch[0].size()o = o.ch[1]default:cnt += o.ch[0].size()if equal {cnt += o.dupCnt}return}}return
}func countRangeSum(nums []int, lower, upper int) (cnt int) {preSum := make([]int, len(nums)+1)for i, v := range nums {preSum[i+1] = preSum[i] + v}t := &treap{}for _, sum := range preSum {left, right := sum-upper, sum-lowercnt += t.rank(right, true) - t.rank(left, false)t.insert(sum)}return
}