的无聊生活

用Common Lisp实现简单的红黑树

发表于 更新于 分类于 6.5k 6 分钟

本来是准备做算法复健，想起来自己曾经也算是正经用 Lisp 工作过，所以试着用 Lisp 写了个简单的红黑树：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217

;; 定义红黑树的节点
(defconstant +red+   'red)
(defconstant +black+ 'black)

(defstruct rbtree
  color
  key
  left
  right
  parent)

;; 插入操作的辅助函数：在二叉搜索树中插入新节点
(defun insert-bst (tree node)
  (let ((parent nil)
        (current tree))
    (loop while current do
         (setf parent current)
         (if (< (rbtree-key node) (rbtree-key current))
             (setf current (rbtree-left current))
             (setf current (rbtree-right current))))
    (setf (rbtree-parent node) parent)
    (if (null parent)
        (setf tree node) ; 树是空的，新节点成为树根
        (if (< (rbtree-key node) (rbtree-key parent))
            (setf (rbtree-left parent) node) ; 成为左孩子
            (setf (rbtree-right parent) node))) ; 成为右孩子
    node))

;; 红黑树的插入操作
(defun insert-rbtree (tree key)
  (let ((new-node (make-rbtree :key key :color +red+)))
    (if (null tree)
        (setf (rbtree-color new-node) +black+)
        (insert-bst tree new-node))
    (fixup-rbtree new-node)
    (root-of new-node))

;; 确保红黑树性质的辅助函数
(defun fixup-rbtree (node)
  (loop while (and (rbtree-parent node)
                   (= +red+ (rbtree-color (rbtree-parent node))))
    (let* ((parent (rbtree-parent node))
           (grandparent (rbtree-parent parent)))
      (if (eq parent (rbtree-left grandparent))
          (let ((uncle (rbtree-right grandparent)))
            (cond ((and uncle (= +red+ (rbtree-color uncle))) ; 情况1
                   (setf (rbtree-color parent) +black+)
                   (setf (rbtree-color uncle) +black+)
                   (setf (rbtree-color grandparent) +red+)
                   (setf node grandparent))
                  ((eq node (rbtree-right parent)) ; 情况2
                   (setf node parent)
                   (left-rotate node))
                  (t ; 情况3
                   (setf (rbtree-color parent) +black+)
                   (setf (rbtree-color grandparent) +red+)
                   (right-rotate grandparent))))
          (let ((uncle (rbtree-left grandparent))) ; 右子树的情况与左子树相反
            (cond ((and uncle (= +red+ (rbtree-color uncle))) ; 情况1
                   (setf (rbtree-color parent) +black+)
                   (setf (rbtree-color uncle) +black+)
                   (setf (rbtree-color grandparent) +red+)
                   (setf node grandparent))
                  ((eq node (rbtree-left parent)) ; 情况2
                   (setf node parent)
                   (right-rotate node))
                  (t ; 情况3
                   (setf (rbtree-color parent) +black+)
                   (setf (rbtree-color grandparent) +red+)
                   (left-rotate grandparent))))))
  (setf (rbtree-color (root-of node)) +black+)) ; 确保根节点是黑色

;; 对节点 x 进行左旋转。
(defun left-rotate (x)
  (let ((y (rbtree-right x)))
    (setf (rbtree-right x) (rbtree-left y)) ; 将 y 的左子树设置为 x 的右子树
    (when (rbtree-left y)
      (setf (rbtree-parent (rbtree-left y)) x)) ; 更新 y 左子树的父节点为 x
    (setf (rbtree-parent y) (rbtree-parent x)) ; 将 x 的父节点设置为 y 的父节点
    (if (null (rbtree-parent x)) ; 如果 x 是根节点
        (setf (rbtree-parent y) nil) ; y 成为新的根节点
        (if (eq x (rbtree-left (rbtree-parent x))) ; 如果 x 是其父节点的左子节点
            (setf (rbtree-left (rbtree-parent x)) y) ; 将 y 设置为 x 的父节点的左子节点
            (setf (rbtree-right (rbtree-parent x)) y))) ; 否则，将 y 设置为 x 的父节点的右子节点
    (setf (rbtree-left y) x) ; 将 x 设置为 y 的左子节点
    (setf (rbtree-parent x) y))) ; 更新 x 的父节点为 y

;; 对节点 y 进行右旋转。
(defun right-rotate (y)
  (let ((x (rbtree-left y)))
    (setf (rbtree-left y) (rbtree-right x)) ; 将 x 的右子树设置为 y 的左子树
    (when (rbtree-right x)
      (setf (rbtree-parent (rbtree-right x)) y)) ; 更新 x 右子树的父节点为 y
    (setf (rbtree-parent x) (rbtree-parent y)) ; 将 y 的父节点设置为 x 的父节点
    (if (null (rbtree-parent y)) ; 如果 y 是根节点
        (setf (rbtree-parent x) nil) ; x 成为新的根节点
        (if (eq y (rbtree-right (rbtree-parent y))) ; 如果 y 是其父节点的右子节点
            (setf (rbtree-right (rbtree-parent y)) x) ; 将 x 设置为 y 的父节点的右子节点
            (setf (rbtree-left (rbtree-parent y)) x))) ; 否则，将 x 设置为 y 的父节点的左子节点
    (setf (rbtree-right x) y) ; 将 y 设置为 x 的右子节点
    (setf (rbtree-parent y) x))) ; 更新 y 的父节点为 x


;; 红黑树的删除操作
(defun transplant (tree u v)
  "用节点 v 替换节点 u。"
  (if (null (rbtree-parent u))
      (setf tree v)
      (if (eq u (rbtree-left (rbtree-parent u)))
          (setf (rbtree-left (rbtree-parent u)) v)
          (setf (rbtree-right (rbtree-parent u)) v)))
  (when v
    (setf (rbtree-parent v) (rbtree-parent u))))

(defun tree-minimum (node)
  "找到以 node 为根的子树的最小节点。"
  (loop while (rbtree-left node) do
       (setf node (rbtree-left node)))
  node)

(defun delete-rbtree (tree key)
  "从红黑树中删除键为 key 的节点。"
  (let ((z (tree-search tree key)) ; 需要实现 tree-search 函数
        (y z)
        (y-original-color (rbtree-color y))
        x)
    (if (null (rbtree-left z))
        (progn
          (setf x (rbtree-right z))
          (transplant tree z (rbtree-right z)))
        (if (null (rbtree-right z))
            (progn
              (setf x (rbtree-left z))
              (transplant tree z (rbtree-left z)))
            (progn
              (setf y (tree-minimum (rbtree-right z)))
              (setf y-original-color (rbtree-color y))
              (setf x (rbtree-right y))
              (if (eq y (rbtree-parent z))
                  (setf (rbtree-parent x) y)
                  (progn
                    (transplant tree y (rbtree-right y))
                    (setf (rbtree-right y) (rbtree-right z))
                    (setf (rbtree-parent (rbtree-right y)) y)))
              (transplant tree z y)
              (setf (rbtree-left y) (rbtree-left z))
              (setf (rbtree-parent (rbtree-left y)) y)
              (setf (rbtree-color y) (rbtree-color z)))))
    (when (eq +black+ y-original-color)
      (fixup-delete-rbtree tree x))
    tree))

;; 调整树以保持红黑性质。
(defun fixup-delete-rbtree (tree x)
  (loop while (and x (not (eq x tree)) (eq +black+ (rbtree-color x)))
    (if (eq x (rbtree-left (rbtree-parent x)))
        (let ((w (rbtree-right (rbtree-parent x))))
          (cond
            ((eq +red+ (rbtree-color w)) ; 情况1：x的兄弟节点w是红色
             (setf (rbtree-color w) +black+)
             (setf (rbtree-color (rbtree-parent x)) +red+)
             (left-rotate (rbtree-parent x))
             (setf w (rbtree-right (rbtree-parent x))))
            ((and (eq +black+ (rbtree-color (rbtree-left w)))
                  (eq +black+ (rbtree-color (rbtree-right w)))) ; 情况2：w的两个子节点都是黑色
             (setf (rbtree-color w) +red+)
             (setf x (rbtree-parent x)))
            (t
              (when (eq +black+ (rbtree-color (rbtree-right w))) ; 情况3：w的右子节点是黑色
                (setf (rbtree-color (rbtree-left w)) +black+)
                (setf (rbtree-color w) +red+)
                (right-rotate w)
                (setf w (rbtree-right (rbtree-parent x))))
              (setf (rbtree-color w) (rbtree-color (rbtree-parent x)))
              (setf (rbtree-color (rbtree-parent x)) +black+)
              (setf (rbtree-color (rbtree-right w)) +black+)
              (left-rotate (rbtree-parent x))
              (setf x tree))))
        (let ((w (rbtree-left (rbtree-parent x)))) ; x 是其父节点的右子节点的情况
          (cond
            ((eq +red+ (rbtree-color w)) ; 情况1：x的兄弟节点w是红色
             (setf (rbtree-color w) +black+)
             (setf (rbtree-color (rbtree-parent x)) +red+)
             (right-rotate tree (rbtree-parent x))
             (setf w (rbtree-left (rbtree-parent x))))
            ((and (eq +black+ (rbtree-color (rbtree-right w)))
                  (eq +black+ (rbtree-color (rbtree-left w)))) ; 情况2：w的两个子节点都是黑色
             (setf (rbtree-color w) +red+)
             (setf x (rbtree-parent x)))
            (t
              (when (eq +black+ (rbtree-color (rbtree-left w))) ; 情况3：w的左子节点是黑色
                (setf (rbtree-color (rbtree-right w)) +black+)
                (setf (rbtree-color w) +red+)
                (left-rotate tree w)
                (setf w (rbtree-left (rbtree-parent x))))
              (setf (rbtree-color w) (rbtree-color (rbtree-parent x)))
              (setf (rbtree-color (rbtree-parent x)) +black+)
              (setf (rbtree-color (rbtree-left w)) +black+)
              (right-rotate tree (rbtree-parent x))
              (setf x tree)))))
    (setf (rbtree-color x) +black+))

;; 在以 node 为根的红黑树中查找键为 key 的节点。
(defun tree-search (node key)
  (loop while (and node (not (eq key (rbtree-key node))))
    do (if (< key (rbtree-key node))
           (setf node (rbtree-left node))
           (setf node (rbtree-right node))))
  node)


;; 找到并返回树的根节点
(defun root-of (node)
  (loop while (rbtree-parent node) do
       (setf node (rbtree-parent node)))
  node)

以上程序写了我整整两天，想当年面试都能直接手写红黑树 …… 其实还有很大的优化空间，同时保不齐还有 Bug。

并且 Common Lisp 的可读性实在太差了。