AVL树/红黑树实现及map与set的封装

Shaw0xyz

1. 引言

在计算机科学中，平衡二叉搜索树（BST）是一种保持树的高度接近最小以确保操作高效的数据结构。常见的平衡二叉搜索树有AVL树和红黑树。本文将详细介绍AVL树和红黑树的实现，并讨论如何将它们封装成map和set数据结构。

1.1 AVL树

AVL树是一种自平衡二叉搜索树，每个节点维护一个平衡因子，即其左右子树的高度差。AVL树通过旋转操作保持平衡，以确保插入、删除和查找操作的时间复杂度为O(log n)。

1.1.1 AVL树的基本操作

AVL树的基本操作包括插入、删除和查找。在插入和删除节点后，AVL树需要进行平衡操作。

(1) 插入操作

插入节点后，如果树失去平衡，需要通过单旋转或双旋转恢复平衡。

示例代码：

1. type AVLNode struct {
   
2.     key    int
   
3.     height int
   
4.     left   *AVLNode
   
5.     right  *AVLNode
   
6. }
   
7. 
   
8. func height(node *AVLNode) int {
   
9.     if node == nil {
   
10.         return 0
    
11.     }
    
12.     return node.height
    
13. }
    
14. 
    
15. func max(a, b int) int {
    
16.     if a > b {
    
17.         return a
    
18.     }
    
19.     return b
    
20. }
    
21. 
    
22. func rightRotate(y *AVLNode) *AVLNode {
    
23.     x := y.left
    
24.     T2 := x.right
    
25. 
    
26.     x.right = y
    
27.     y.left = T2
    
28. 
    
29.     y.height = max(height(y.left), height(y.right)) + 1
    
30.     x.height = max(height(x.left), height(x.right)) + 1
    
31. 
    
32.     return x
    
33. }
    
34. 
    
35. func leftRotate(x *AVLNode) *AVLNode {
    
36.     y := x.right
    
37.     T2 := y.left
    
38. 
    
39.     y.left = x
    
40.     x.right = T2
    
41. 
    
42.     x.height = max(height(x.left), height(x.right)) + 1
    
43.     y.height = max(height(y.left), height(y.right)) + 1
    
44. 
    
45.     return y
    
46. }
    
47. 
    
48. func getBalance(node *AVLNode) int {
    
49.     if node == nil {
    
50.         return 0
    
51.     }
    
52.     return height(node.left) - height(node.right)
    
53. }
    
54. 
    
55. func insert(node *AVLNode, key int) *AVLNode {
    
56.     if node == nil {
    
57.         return &AVLNode{key: key, height: 1}
    
58.     }
    
59.     if key < node.key {
    
60.         node.left = insert(node.left, key)
    
61.     } else if key > node.key {
    
62.         node.right = insert(node.right, key)
    
63.     } else {
    
64.         return node
    
65.     }
    
66. 
    
67.     node.height = 1 + max(height(node.left), height(node.right))
    
68. 
    
69.     balance := getBalance(node)
    
70. 
    
71.     if balance > 1 && key < node.left.key {
    
72.         return rightRotate(node)
    
73.     }
    
74. 
    
75.     if balance < -1 && key > node.right.key {
    
76.         return leftRotate(node)
    
77.     }
    
78. 
    
79.     if balance > 1 && key > node.left.key {
    
80.         node.left = leftRotate(node.left)
    
81.         return rightRotate(node)
    
82.     }
    
83. 
    
84.     if balance < -1 && key < node.right.key {
    
85.         node.right = rightRotate(node.right)
    
86.         return leftRotate(node)
    
87.     }
    
88. 
    
89.     return node
    
90. }

复制代码

(2) 删除操作

删除节点后，同样需要进行旋转操作以保持平衡。

1.1.2 AVL树的查找操作

查找操作与普通二叉搜索树相同，通过比较节点值递归或迭代进行查找。

示例代码：

1. func search(node *AVLNode, key int) *AVLNode {
   
2.     if node == nil || node.key == key {
   
3.         return node
   
4.     }
   
5.     if key < node.key {
   
6.         return search(node.left, key)
   
7.     }
   
8.     return search(node.right, key)
   
9. }

复制代码

1.2 红黑树

红黑树是一种自平衡二叉搜索树，确保在插入和删除操作后保持近似平衡。红黑树通过颜色属性和旋转操作来保持平衡，插入和删除的时间复杂度为O(log n)。

1.2.1 红黑树的基本操作

红黑树的基本操作包括插入、删除和查找。

(1) 插入操作

插入节点后，通过重新着色和旋转操作恢复红黑树的性质。

(2) 删除操作

删除节点后，通过重新着色和旋转操作保持平衡。

1.2.2 红黑树的查找操作

查找操作与普通二叉搜索树相同，通过比较节点值递归或迭代进行查找。

1.3 map和set的封装

通过AVL树或红黑树实现map和set数据结构，可以确保键的有序性和操作的高效性。

1.3.1 map的实现

map是一种键值对数据结构，可以通过AVL树或红黑树实现。每个节点存储一个键值对，支持插入、删除和查找操作。

示例代码：

1. type MapNode struct {
   
2.     key    int
   
3.     value  interface{}
   
4.     height int
   
5.     left   *MapNode
   
6.     right  *MapNode
   
7. }
   
8. 
   
9. func insertMapNode(node *MapNode, key int, value interface{}) *MapNode {
   
10.     if node == nil {
    
11.         return &MapNode{key: key, value: value, height: 1}
    
12.     }
    
13.     if key < node.key {
    
14.         node.left = insertMapNode(node.left, key, value)
    
15.     } else if key > node.key {
    
16.         node.right = insertMapNode(node.right, key, value)
    
17.     } else {
    
18.         node.value = value
    
19.         return node
    
20.     }
    
21. 
    
22.     node.height = 1 + max(height(node.left), height(node.right))
    
23. 
    
24.     balance := getBalance(node)
    
25. 
    
26.     if balance > 1 && key < node.left.key {
    
27.         return rightRotate(node)
    
28.     }
    
29. 
    
30.     if balance < -1 && key > node.right.key {
    
31.         return leftRotate(node)
    
32.     }
    
33. 
    
34.     if balance > 1 && key > node.left.key {
    
35.         node.left = leftRotate(node.left)
    
36.         return rightRotate(node)
    
37.     }
    
38. 
    
39.     if balance < -1 && key < node.right.key {
    
40.         node.right = rightRotate(node.right)
    
41.         return leftRotate(node)
    
42.     }
    
43. 
    
44.     return node
    
45. }

复制代码

1.3.2 set的实现

set是一种无重复元素的数据结构，可以通过AVL树或红黑树实现。每个节点存储一个元素，支持插入、删除和查找操作。

示例代码：

1. type SetNode struct {
   
2.     key    int
   
3.     height int
   
4.     left   *SetNode
   
5.     right  *SetNode
   
6. }
   
7. 
   
8. func insertSetNode(node *SetNode, key int) *SetNode {
   
9.     if node == nil {
   
10.         return &SetNode{key: key, height: 1}
    
11.     }
    
12.     if key < node.key {
    
13.         node.left = insertSetNode(node.left, key)
    
14.     } else if key > node.key {
    
15.         node.right = insertSetNode(node.right, key)
    
16.     } else {
    
17.         return node
    
18.     }
    
19. 
    
20.     node.height = 1 + max(height(node.left), height(node.right))
    
21. 
    
22.     balance := getBalance(node)
    
23. 
    
24.     if balance > 1 && key < node.left.key {
    
25.         return rightRotate(node)
    
26.     }
    
27. 
    
28.     if balance < -1 && key > node.right.key {
    
29.         return leftRotate(node)
    
30.     }
    
31. 
    
32.     if balance > 1 && key > node.left.key {
    
33.         node.left = leftRotate(node.left)
    
34.         return rightRotate(node)
    
35.     }
    
36. 
    
37.     if balance < -1 && key < node.right.key {
    
38.         node.right = rightRotate(node.right)
    
39.         return leftRotate(node)
    
40.     }
    
41. 
    
42.     return node
    
43. }

复制代码

2. 结论

通过本文的介绍，我们了解了AVL树和红黑树的基本原理及实现方法，并讨论了如何将它们封装成map和set数据结构。掌握这些高级数据结构和算法对于提高程序性能和解决复杂问题至关重要。希望这篇文章能对你理解和使用平衡二叉搜索树有所帮助。






/ 荔枝学姐de课后专栏 /

Hi！这里是荔枝学姐~

欢迎来到我的课后专栏

自然语言学渣 NLP摆烂姐

热衷于技术写作 IT边角料

AIGC & Coding & linux ...

~互撩~ TG: @Shaw_0xyz