From 5776e0ba6e59563c0efd1b397e35a7e59de627f6 Mon Sep 17 00:00:00 2001 From: bandl <1658002533@qq.com> Date: Tue, 12 Oct 2021 15:10:00 +0800 Subject: [PATCH] feat(util): add skiplist --- pkg/util/skiplist/skiplist.go | 307 ++++++++++++++++++++++++++++++++++ 1 file changed, 307 insertions(+) create mode 100644 pkg/util/skiplist/skiplist.go diff --git a/pkg/util/skiplist/skiplist.go b/pkg/util/skiplist/skiplist.go new file mode 100644 index 0000000..391f559 --- /dev/null +++ b/pkg/util/skiplist/skiplist.go @@ -0,0 +1,307 @@ +package skiplist + +import ( + "crypto/rand" + "fmt" + "math" + "math/big" +) + +type skipListNode struct { + score float64 + val interface{} + next *skipListNode + down *skipListNode + up *skipListNode + pre *skipListNode +} + +func newSkipListNode(score float64, val interface{}) *skipListNode { + if val == nil { + return nil + } + + return &skipListNode{ + score: score, + val: val, + } +} + +type SkipList struct { + head *skipListNode + tail *skipListNode + length int + levels int // 单前层级 + + maxLevels int // 最大层级 +} + +func NewSkipList(maxLevel int) *SkipList { + skl := new(SkipList) + skl.head = &skipListNode{score: math.MinInt64} + skl.tail = &skipListNode{score: math.MaxInt64} + + skl.head.next = skl.tail + skl.tail.pre = skl.head + + skl.length = 0 + skl.levels = 1 + skl.maxLevels = maxLevel + + return skl +} + +func (s *SkipList) Length() int { + return s.length +} + +func (s *SkipList) Levels() int { + return s.levels +} + +func (s *SkipList) MaxLength() int { + return s.maxLevels +} + +// Insert 插入数据 +func (s *SkipList) Insert(score float64, val interface{}) { + + node := newSkipListNode(score, val) + f := s.searchNode(score) + + s.insertAfter(f, node) + + // 随机上升 + + for newLevel := 1; ; newLevel++ { + + rander, _ := rand.Int(rand.Reader, big.NewInt(2)) + + if rander.Int64() == 0 || newLevel >= s.maxLevels { + break + } + + // 上升层级 + if newLevel >= s.levels && s.levels < s.maxLevels { + s.newLevels() + } + + for f.up == nil { + f = f.pre + } + + f = f.up + + tmpNode := &skipListNode{score: score} + node.up = tmpNode + tmpNode.down = node + s.insertAfter(f, tmpNode) + + node = tmpNode + } + + s.length += 1 +} + +// Pop 弹出随机一个 score 值的 val +func (s *SkipList) Pop(score float64) interface{} { + f := s.searchNode(score) + if f.score == score { + return nil + } + v := f.val + + s.delSkipListNode(f) + return v +} + +// Get 获取 随机一个 score 值的 val +func (s *SkipList) Get(score float64) interface{} { + node := s.searchNode(score) + if node != nil && node.score == score { + return node.val + } + return nil +} + +// GetAll 获取全部的 ALL +func (s *SkipList) GetAll(score float64) []interface{} { + node := s.searchNode(score) + values := make([]interface{}, 0) + p := node + + // pre + for p.score == score { + values = append(values, p.val) + p = p.pre + } + + // next + p = node.next + for p.score == score { + values = append(values, p.val) + p = p.next + } + + return values +} + +// RemoveAll 删除全部 +func (s *SkipList) RemoveAll(score float64) { + node := s.searchNode(score) + p := node + delNode := make([]*skipListNode, 0) + // pre + for p.score == score { + delNode = append(delNode, p) + p = p.pre + } + + // next + p = node.next + for p.score == score { + delNode = append(delNode, p) + p = p.next + } + s.delSkipListNode(delNode...) +} + +// ClearLeft 删除 小等于 score 的全部节点 +func (s *SkipList) ClearLeft(score float64) *skipListNode { + head := s.head + + for head.down != nil { + p := head + for p.score <= score { + p = p.next + } + + head.next = p + p.pre = head + head = head.down + } + + node := s.searchNode(score) + for node.score <= score { + node = node.next + } + + oldNode := node.pre + head.next = node + + node.pre = head + return oldNode +} + +func (s *SkipList) PopLeft(score float64) []interface{} { + node := s.ClearLeft(score) + values := make([]interface{}, 0, 10) + for node.score != math.MinInt64 { + values = append(values, node.val) + node = node.pre + } + return values +} + +// 查找节点位置 +func (s *SkipList) searchNode(score float64) *skipListNode { + p := s.head + for p != nil { + // 尾部是极大极小值 + if score == p.score { + if p.down == nil { + return p + } + p = p.down + continue + } + + if score > p.score && score < p.next.score { + if p.down == nil { + return p + } + p = p.down + continue + } + + p = p.next + } + + return p +} + +// 在节点后插入新节点 +func (s *SkipList) insertAfter(pNode *skipListNode, newNode *skipListNode) { + newNode.next = pNode.next + newNode.pre = pNode + pNode.next.pre = newNode + pNode.next = newNode +} + +// 添加 一个新的 levels 层 +func (s *SkipList) newLevels() { + newHead := &skipListNode{score: math.MinInt64} + newTail := &skipListNode{score: math.MaxInt64} + + newHead.next = newTail + newTail.pre = newHead + + s.head.up = newHead + newHead.down = s.head + s.tail.up = newTail + newTail.down = s.tail + + s.head = newHead + s.tail = newTail + s.levels++ +} + +func (s *SkipList) debugPrint() { + mapScore := make(map[float64]int) + p := s.head + for p.down != nil { + p = p.down + } + index := 0 + for p != nil { + mapScore[p.score] = index + p = p.next + index++ + } + p = s.head + for i := 0; i < s.levels; i++ { + q := p + preIndex := 0 + for q != nil { + s := q.score + if s == math.MinInt64 { + fmt.Printf("%s", "BEGIN") + q = q.next + continue + } + index := mapScore[s] + c := (index - preIndex - 1) * 6 + for m := 0; m < c; m++ { + fmt.Print("-") + } + if s == math.MaxInt64 { + fmt.Printf("-->%s\n", "END") + } else { + fmt.Printf("-->%3d", int(s)) + preIndex = index + } + q = q.next + } + p = p.down + } +} + +func (s *SkipList) delSkipListNode(sKm ...*skipListNode) { + for _, f := range sKm { + for f != nil { + f.pre.next = f.next + f.next.pre = f.pre + f = f.up + } + } +}