xmengnet/hello-algo
1
0
Fork 0
mirror of https://github.com/krahets/hello-algo.git synced 2024-12-26 13:36:29 +08:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge pull request #224 from Reanon/feat/go-avl-tree

Browse source 
feat(tree/avl_tree): add go code
This commit is contained in:
Yudong Jin 2023-01-09 02:04:45 +08:00 committed by GitHub
commit e5bc9436fa
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 433 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

211
codes/go/chapter_tree/avl_tree.go Normal file
View file

@ -0,0 +1,211 @@
// File: avl_tree.go
// Created Time: 2023-01-08
// Author: Reanon (793584285@qq.com)
package chapter_tree
import . "github.com/krahets/hello-algo/pkg"
/* AVL Tree*/
type avlTree struct {
// 根节点
root *TreeNode
}
func newAVLTree() *avlTree {
return &avlTree{root: nil}
}
/* 获取结点高度 */
func height(node *TreeNode) int {
// 空结点高度为 -1 ，叶结点高度为 0
if node != nil {
return node.Height
}
return -1
}
/* 更新结点高度 */
func updateHeight(node *TreeNode) {
lh := height(node.Left)
rh := height(node.Right)
// 结点高度等于最高子树高度 + 1
if lh > rh {
node.Height = lh + 1
} else {
node.Height = rh + 1
}
}
/* 获取平衡因子 */
func balanceFactor(node *TreeNode) int {
// 空结点平衡因子为 0
if node == nil {
return 0
}
// 结点平衡因子 = 左子树高度 - 右子树高度
return height(node.Left) - height(node.Right)
}
/* 右旋操作 */
func rightRotate(node *TreeNode) *TreeNode {
child := node.Left
grandChild := child.Right
// 以 child 为原点，将 node 向右旋转
child.Right = node
node.Left = grandChild
// 更新结点高度
updateHeight(node)
updateHeight(child)
// 返回旋转后子树的根节点
return child
}
/* 左旋操作 */
func leftRotate(node *TreeNode) *TreeNode {
child := node.Right
grandChild := child.Left
// 以 child 为原点，将 node 向左旋转
child.Left = node
node.Right = grandChild
// 更新结点高度
updateHeight(node)
updateHeight(child)
// 返回旋转后子树的根节点
return child
}
/* 执行旋转操作，使该子树重新恢复平衡 */
func rotate(node *TreeNode) *TreeNode {
// 获取结点 node 的平衡因子
// Go 推荐短变量，这里 bf 指代 balanceFactor
bf := balanceFactor(node)
// 左偏树
if bf > 1 {
if balanceFactor(node.Left) >= 0 {
// 右旋
return rightRotate(node)
} else {
// 先左旋后右旋
node.Left = leftRotate(node.Left)
return rightRotate(node)
}
}
// 右偏树
if bf < -1 {
if balanceFactor(node.Right) <= 0 {
// 左旋
return leftRotate(node)
} else {
// 先右旋后左旋
node.Right = rightRotate(node.Right)
return leftRotate(node)
}
}
// 平衡树，无需旋转，直接返回
return node
}
/* 插入结点 */
func (t *avlTree) insert(val int) *TreeNode {
t.root = insertHelper(t.root, val)
return t.root
}
/* 递归插入结点（辅助函数） */
func insertHelper(node *TreeNode, val int) *TreeNode {
if node == nil {
return NewTreeNode(val)
}
/* 1. 查找插入位置，并插入结点 */
if val < node.Val {
node.Left = insertHelper(node.Left, val)
} else if val > node.Val {
node.Right = insertHelper(node.Right, val)
} else {
// 重复结点不插入，直接返回
return node
}
// 更新结点高度
updateHeight(node)
/* 2. 执行旋转操作，使该子树重新恢复平衡 */
node = rotate(node)
// 返回子树的根节点
return node
}
/* 删除结点 */
func (t *avlTree) remove(val int) *TreeNode {
root := removeHelper(t.root, val)
return root
}
/* 递归删除结点（辅助函数） */
func removeHelper(node *TreeNode, val int) *TreeNode {
if node == nil {
return nil
}
/* 1. 查找结点，并删除之 */
if val < node.Val {
node.Left = removeHelper(node.Left, val)
} else if val > node.Val {
node.Right = removeHelper(node.Right, val)
} else {
if node.Left == nil || node.Right == nil {
child := node.Left
if node.Right != nil {
child = node.Right
}
// 子结点数量 = 0 ，直接删除 node 并返回
if child == nil {
return nil
} else {
// 子结点数量 = 1 ，直接删除 node
node = child
}
} else {
// 子结点数量 = 2 ，则将中序遍历的下个结点删除，并用该结点替换当前结点
temp := getInOrderNext(node.Right)
node.Right = removeHelper(node.Right, temp.Val)
node.Val = temp.Val
}
}
// 更新结点高度
updateHeight(node)
/* 2. 执行旋转操作，使该子树重新恢复平衡 */
node = rotate(node)
// 返回子树的根节点
return node
}
/* 获取中序遍历中的下一个结点（仅适用于 root 有左子结点的情况） */
func getInOrderNext(node *TreeNode) *TreeNode {
if node == nil {
return node
}
// 循环访问左子结点，直到叶结点时为最小结点，跳出
for node.Left != nil {
node = node.Left
}
return node
}
/* 查找结点 */
func (t *avlTree) search(val int) *TreeNode {
cur := t.root
// 循环查找，越过叶结点后跳出
for cur != nil {
// 目标结点在 root 的右子树中
if cur.Val < val {
cur = cur.Right
} else if cur.Val > val {
// 目标结点在 root 的左子树中
cur = cur.Left
} else {
// 找到目标结点，跳出循环
break
}
}
// 返回目标结点
return cur
}

54
codes/go/chapter_tree/avl_tree_test.go Normal file
View file

@ -0,0 +1,54 @@
// File: avl_tree_test.go
// Created Time: 2023-01-08
// Author: Reanon (793584285@qq.com)
package chapter_tree
import (
"fmt"
"testing"
. "github.com/krahets/hello-algo/pkg"
)
func TestAVLTree(t *testing.T) {
/* 初始化空 AVL 树 */
tree := newAVLTree()
/* 插入结点 */
// 请关注插入结点后，AVL 树是如何保持平衡的
testInsert(tree, 1)
testInsert(tree, 2)
testInsert(tree, 3)
testInsert(tree, 4)
testInsert(tree, 5)
testInsert(tree, 8)
testInsert(tree, 7)
testInsert(tree, 9)
testInsert(tree, 10)
testInsert(tree, 6)
/* 插入重复结点 */
testInsert(tree, 7)
/* 删除结点 */
// 请关注删除结点后，AVL 树是如何保持平衡的
testRemove(tree, 8) // 删除度为 0 的结点
testRemove(tree, 5) // 删除度为 1 的结点
testRemove(tree, 4) // 删除度为 2 的结点
/* 查询结点 */
node := tree.search(7)
fmt.Printf("\n查找到的结点对象为 %#v ，结点值 = %d \n", node, node.Val)
}
func testInsert(tree *avlTree, val int) {
tree.insert(val)
fmt.Printf("\n插入结点 %d 后，AVL 树为 \n", val)
PrintTree(tree.root)
}
func testRemove(tree *avlTree, val int) {
tree.remove(val)
fmt.Printf("\n删除结点 %d 后，AVL 树为 \n", val)
PrintTree(tree.root)
}

14
codes/go/pkg/tree_node.go
View file

@ -9,16 +9,18 @@ import (
) )
type TreeNode struct { type TreeNode struct {
Val int Val int // 结点值
Left *TreeNode Height int // 结点高度
Right *TreeNode Left *TreeNode // 左子结点引用
Right *TreeNode // 右子结点引用
} }
func NewTreeNode(v int) *TreeNode { func NewTreeNode(v int) *TreeNode {
return &TreeNode{ return &TreeNode{
Left: nil, Val: v,
Right: nil, Height: 0,
Val: v, Left: nil,
Right: nil,
} }
} }

165
docs/chapter_tree/avl_tree.md
View file

@ -60,7 +60,13 @@ G. M. Adelson-Velsky 和 E. M. Landis 在其 1962 年发表的论文 "An algorit
=== "Go" === "Go"
```go title="avl_tree.go" ```go title="avl_tree.go"
/* AVL 树结点类 */
type TreeNode struct {
Val int // 结点值
Height int // 结点高度
Left *TreeNode // 左子结点引用
Right *TreeNode // 右子结点引用
}
``` ```
=== "JavaScript" === "JavaScript"
@ -143,7 +149,26 @@ G. M. Adelson-Velsky 和 E. M. Landis 在其 1962 年发表的论文 "An algorit
=== "Go" === "Go"
```go title="avl_tree.go" ```go title="avl_tree.go"
/* 获取结点高度 */
func height(node *TreeNode) int {
// 空结点高度为 -1 ，叶结点高度为 0
if node != nil {
return node.Height
}
return -1
}
/* 更新结点高度 */
func updateHeight(node *TreeNode) {
lh := height(node.Left)
rh := height(node.Right)
// 结点高度等于最高子树高度 + 1
if lh > rh {
node.Height = lh + 1
} else {
node.Height = rh + 1
}
}
``` ```
=== "JavaScript" === "JavaScript"
@ -225,7 +250,15 @@ G. M. Adelson-Velsky 和 E. M. Landis 在其 1962 年发表的论文 "An algorit
=== "Go" === "Go"
```go title="avl_tree.go" ```go title="avl_tree.go"
/* 获取平衡因子 */
func balanceFactor(node *TreeNode) int {
// 空结点平衡因子为 0
if node == nil {
return 0
}
// 结点平衡因子 = 左子树高度 - 右子树高度
return height(node.Left) - height(node.Right)
}
``` ```
=== "JavaScript" === "JavaScript"
@ -338,7 +371,19 @@ AVL 树的独特之处在于「旋转 Rotation」的操作，其可 **在不影
=== "Go" === "Go"
```go title="avl_tree.go" ```go title="avl_tree.go"
/* 右旋操作 */
func rightRotate(node *TreeNode) *TreeNode {
child := node.Left
grandChild := child.Right
// 以 child 为原点，将 node 向右旋转
child.Right = node
node.Left = grandChild
// 更新结点高度
updateHeight(node)
updateHeight(child)
// 返回旋转后子树的根节点
return child
}
``` ```
=== "JavaScript" === "JavaScript"
@ -441,7 +486,19 @@ AVL 树的独特之处在于「旋转 Rotation」的操作，其可 **在不影
=== "Go" === "Go"
```go title="avl_tree.go" ```go title="avl_tree.go"
/* 左旋操作 */
func leftRotate(node *TreeNode) *TreeNode {
child := node.Right
grandChild := child.Left
// 以 child 为原点，将 node 向左旋转
child.Left = node
node.Right = grandChild
// 更新结点高度
updateHeight(node)
updateHeight(child)
// 返回旋转后子树的根节点
return child
}
``` ```
=== "JavaScript" === "JavaScript"
@ -592,7 +649,36 @@ AVL 树的独特之处在于「旋转 Rotation」的操作，其可 **在不影
=== "Go" === "Go"
```go title="avl_tree.go" ```go title="avl_tree.go"
/* 执行旋转操作，使该子树重新恢复平衡 */
func rotate(node *TreeNode) *TreeNode {
// 获取结点 node 的平衡因子
// Go 推荐短变量，这里 bf 指代 balanceFactor
bf := balanceFactor(node)
// 左偏树
if bf > 1 {
if balanceFactor(node.Left) >= 0 {
// 右旋
return rightRotate(node)
} else {
// 先左旋后右旋
node.Left = leftRotate(node.Left)
return rightRotate(node)
}
}
// 右偏树
if bf < -1 {
if balanceFactor(node.Right) <= 0 {
// 左旋
return leftRotate(node)
} else {
// 先右旋后左旋
node.Right = rightRotate(node.Right)
return leftRotate(node)
}
}
// 平衡树，无需旋转，直接返回
return node
}
``` ```
=== "JavaScript" === "JavaScript"
@ -730,7 +816,32 @@ AVL 树的独特之处在于「旋转 Rotation」的操作，其可 **在不影
=== "Go" === "Go"
```go title="avl_tree.go" ```go title="avl_tree.go"
/* 插入结点 */
func (t *avlTree) insert(val int) *TreeNode {
t.root = insertHelper(t.root, val)
return t.root
}
/* 递归插入结点（辅助函数） */
func insertHelper(node *TreeNode, val int) *TreeNode {
if node == nil {
return NewTreeNode(val)
}
/* 1. 查找插入位置，并插入结点 */
if val < node.Val {
node.Left = insertHelper(node.Left, val)
} else if val > node.Val {
node.Right = insertHelper(node.Right, val)
} else {
// 重复结点不插入，直接返回
return node
}
// 更新结点高度
updateHeight(node)
/* 2. 执行旋转操作，使该子树重新恢复平衡 */
node = rotate(node)
// 返回子树的根节点
return node
}
``` ```
=== "JavaScript" === "JavaScript"
@ -876,7 +987,49 @@ AVL 树的独特之处在于「旋转 Rotation」的操作，其可 **在不影
=== "Go" === "Go"
```go title="avl_tree.go" ```go title="avl_tree.go"
/* 删除结点 */
func (t *avlTree) remove(val int) *TreeNode {
root := removeHelper(t.root, val)
return root
}
/* 递归删除结点（辅助函数） */
func removeHelper(node *TreeNode, val int) *TreeNode {
if node == nil {
return nil
}
/* 1. 查找结点，并删除之 */
if val < node.Val {
node.Left = removeHelper(node.Left, val)
} else if val > node.Val {
node.Right = removeHelper(node.Right, val)
} else {
if node.Left == nil || node.Right == nil {
child := node.Left
if node.Right != nil {
child = node.Right
}
// 子结点数量 = 0 ，直接删除 node 并返回
if child == nil {
return nil
} else {
// 子结点数量 = 1 ，直接删除 node
node = child
}
} else {
// 子结点数量 = 2 ，则将中序遍历的下个结点删除，并用该结点替换当前结点
temp := getInOrderNext(node.Right)
node.Right = removeHelper(node.Right, temp.Val)
node.Val = temp.Val
}
}
// 更新结点高度
updateHeight(node)
/* 2. 执行旋转操作，使该子树重新恢复平衡 */
node = rotate(node)
// 返回子树的根节点
return node
}
``` ```
=== "JavaScript" === "JavaScript"