1
package skiplist
2

3
import (
4
    "math/rand"
5
)
6

7
const maxLevel = 16 // 最大层数，2^16 个节点足够
8

9
// 节点
10
type node struct {
11
    key   int
12
    value any
13
    next  []*node // 每层的前向指针
14
}
15

16
// 跳表
17
type SkipList struct {
18
    header *node
19
    level  int // 当前最大层数
20
    size   int // 节点数
21
}
22

23
func New() *SkipList {
24
    return &SkipList{
25
        header: &node{next: make([]*node, maxLevel)},
26
        level:  1,
27
    }
28
}
29

30
// 随机层数，p = 0.5
31
func randomLevel() int {
32
    lv := 1
33
    for rand.Float64() < 0.5 && lv < maxLevel {
34
        lv++
35
    }
36
    return lv
37
}
38

39
// 查找
40
func (sl *SkipList) Get(key int) (any, bool) {
41
    cur := sl.header
42
    // 从最高层往下找
43
    for i := sl.level - 1; i >= 0; i-- {
44
        for cur.next[i] != nil && cur.next[i].key < key {
45
            cur = cur.next[i]
46
        }
47
    }
48
    cur = cur.next[0] // 移到第 0 层的下一个节点
49
    if cur != nil && cur.key == key {
50
        return cur.value, true
51
    }
52
    return nil, false
53
}
54

55
// 插入
56
func (sl *SkipList) Put(key int, value any) {
57
    // 记录每层的前驱节点
58
    prev := make([]*node, maxLevel)
59
    cur := sl.header
60
    for i := sl.level - 1; i >= 0; i-- {
61
        for cur.next[i] != nil && cur.next[i].key < key {
62
            cur = cur.next[i]
63
        }
64
        prev[i] = cur
65
    }
66

67
    // key 已存在则更新
68
    if nxt := cur.next[0]; nxt != nil && nxt.key == key {
69
        nxt.value = value
70
        return
71
    }
72

73
    lv := randomLevel()
74
    if lv > sl.level {
75
        // 新层数高于当前层数，header 作为新层的前驱
76
        for i := sl.level; i < lv; i++ {
77
            prev[i] = sl.header
78
        }
79
        sl.level = lv
80
    }
81

82
    nd := &node{key: key, value: value, next: make([]*node, lv)}
83
    for i := 0; i < lv; i++ {
84
        nd.next[i] = prev[i].next[i] // 新节点指向后继
85
        prev[i].next[i] = nd         // 前驱指向新节点
86
    }
87
    sl.size++
88
}