手撸golang 基本数据结构与算法 二叉查找树

缘起

最近阅读<<我的第一本算法书>>(【日】石田保辉；宫崎修一)
本系列笔记拟采用golang练习之

二叉查找树

二叉查找树（又叫作二叉搜索树或二叉排序树）是一种数据结构，
数据存储于二叉查找树的各个结点中。

二叉查找树有两个性质:
第一个是每个结点的值均大于其左子树上任意一个结点的值,
第二个是每个结点的值均小于其右子树上任意一个结点的值。

根据这两个性质可以得到以下结论。
首先，二叉查找树的最小结点要从顶端开始，往其左下的末端寻找。
反过来，二叉查找树的最大结点要从顶端开始，往其右下的末端寻找。

摘自 <<我的第一本算法书>> 【日】石田保辉；宫崎修一

目标

• 实现一棵二叉查找树, 并测试其基本功能

设计

• IComparator: 定义值比较函数. 值比较函数可以返回小于, 等于, 大于三种情况
• IBinarySearchTree: 定义二叉查找树的接口, 增删改查都要.
• IIterator: 定义二叉查找树的遍历接口.
• tComparator: 值比较函数的包装器, 实现IComparator接口. 具体比较函数由外部传入.
• tBinarySearchTree: 二叉查找树的示例, 实现IBinarySearchTree接口.
• tBinaryNode: 二叉查找树节点
• tBSTreeIterator: 二叉查找树的遍历迭代器, 内部使用广度优先搜索+候选队列.

单元测试

bstree_test.go, 测试二叉查找树的增删改查, 以及大小序输出.

package data_structure

import (
    "fmt"
    bstree "learning/gooop/data_structure/binary_search_tree"
    "math/rand"
    "strings"
    "testing"
    "time"
)

func Test_BinarySearchTree(t *testing.T) {
    fnAssertTrue := func(b bool, msg string) {
        if !b {
            panic(msg)
        }
    }

    fnCompare := func(a interface{}, b interface{}) bstree.CompareResult {
        i1 := a.(int)
        i2 := b.(int)

        if i1 < i2 {
            return bstree.LESS
        } else if i1 == i2 {
            return bstree.EQUAL
        } else {
            return bstree.GREATER
        }
    }
    comparator := bstree.NewComparator(fnCompare)

    // test empty tree
    tree := bstree.NewBinarySearchTree(comparator)
    t.Log(tree)
    fnAssertTrue(tree.String() == "r=nil,s=0,v=0", "expecting r=nil,s=0,v=0")
    fnAssertTrue(tree.Size() == 0, "expecting size == 0")
    fnAssertTrue(tree.IsEmpty(), "expecting empty")

    // test one item
    tree.Push(5)
    t.Log(tree)
    fnAssertTrue(tree.String() == "r=5,s=1,v=1 5(n,n)", "expecting r=5,s=1,v=1 5(n,n)")
    fnAssertTrue(tree.Size() == 1, "expecting size == 1")
    fnAssertTrue(tree.IsNotEmpty(), "expecting not empty")

    // test min
    ok, v := tree.Min()
    fnAssertTrue(ok, "expecting min() ok")
    fnAssertTrue(v == 5, "expecting 5")

    // test max
    ok, v = tree.Max()
    fnAssertTrue(ok, "expecting max() ok")
    fnAssertTrue(v == 5, "expecting 5")
    fnAssertTrue(tree.String() == "r=5,s=1,v=1 5(n,n)", "expecting r=5,s=1,v=1 5(n,n)")

    // test pop one
    ok, v = tree.PopMin()
    t.Log(tree)
    fnAssertTrue(tree.String() == "r=nil,s=0,v=2", "expecting r=nil,s=0,v=2")
    fnAssertTrue(ok, "expecting PopMin() ok")
    fnAssertTrue(v == 5, "expecting 5")
    fnAssertTrue(tree.Size() == 0, "expecting size == 0")
    fnAssertTrue(tree.IsEmpty(), "expecting empty")

    // test batch push
    samples := []int{ 5,4,8, 2, 7, 9, 1,3,6 }
    for i := 0;i < len(samples);i++ {
        tree.Push(samples[i])
    }
    t.Log(tree)
    fnAssertTrue(tree.Size() == len(samples), "expecting Size() == len(samples)")
    fnAssertTrue(tree.String() == "r=5,s=9,v=11 5(4,8),4(2,n),8(7,9),2(1,3),7(6,n),9(n,n),1(n,n),3(n,n),6(n,n)", "expecting r=5,s=9,v=11 5(4,8),4(2,n),8(7,9),2(1,3),7(6,n),9(n,n),1(n,n),3(n,n),6(n,n)")

    for _,it := range samples {
        fnAssertTrue(tree.Has(it), "expecting Has() == true")
    }

    // test iterator
    iter := tree.Iterator()
    fnAssertTrue(iter.More(), "expectiong More()")
    iterItems := make([]string, 0)
    for range samples {
        ok,v = iter.Next()
        t.Logf("ok=%v, v=%v", true, v)
        fnAssertTrue(ok, "expectiong Next()")
        iterItems = append(iterItems, fmt.Sprintf("%v", v))
    }
    fnAssertTrue(strings.Join(iterItems, ",") == "5,4,8,2,7,9,1,3,6", "expecting 5,4,8,2,7,9,1,3,6")
    fnAssertTrue(iter.More() == false, "expectiong !iter.More()")
    ok,v = iter.Next()
    fnAssertTrue(!ok, "expecting !iter.Next()")

    // test min
    ok,v = tree.Min()
    fnAssertTrue(ok, "expecting ok")
    fnAssertTrue(v == 1, "expection Min() == 1")

    // test max
    ok,v = tree.Max()
    fnAssertTrue(ok, "expecting ok")
    fnAssertTrue(v == 9, "expection Max() == 9")

    // test batch pop min
    for i := 1;i <= 9;i ++ {
        ok,v = tree.PopMin()
        t.Logf("i=%v v=%v tree=%s", i, v, tree.String())
        fnAssertTrue(ok, "expecting ok")
        fnAssertTrue(v == i, fmt.Sprintf("expecting %v", i))
    }

    // test batch pop max
    for i := 0;i < len(samples);i++ {
        tree.Push(samples[i])
    }
    t.Log(tree)
    for i := 1;i <= 9;i ++ {
        ok,v = tree.PopMax()
        t.Logf("i=%v v=%v tree=%s", i, v, tree.String())
        fnAssertTrue(ok, "expecting ok")
        fnAssertTrue(v == 10 - i, fmt.Sprintf("expecting %v", 10 - i))
    }

    // test batch push
    samples = []int{ 2,1,3 }
    for i := 0;i < len(samples);i++ {
        tree.Push(samples[i])
    }
    t.Log(tree)
    fnAssertTrue(tree.String() == "r=2,s=3,v=41 2(1,3),1(n,n),3(n,n)", "expecting 2(1,3),1(n,n),3(n,n)")

    // test clear
    tree.Clear()
    t.Log(tree)
    fnAssertTrue(tree.String() == "r=nil,s=0,v=42", "expecting empty")

    // test batch remove
    rnd := rand.New(rand.NewSource(time.Now().UnixNano()))
    rndItems := make(map[int]bool, 0)
    for i := 0;i < 10000;i++ {
        x := rnd.Intn(50000)
        tree.Push(x)
        rndItems[x] = true
    }
    for k,_ := range rndItems {
        //t.Logf("removing %v", k)
        fnAssertTrue(tree.Remove(k), "expecting Remove() = true")
        //t.Log(tree)
    }
    t.Log(tree)
    fnAssertTrue(tree.IsEmpty(), "expecting empty")
}

测试输出

$ go test -v bstree_test.go 
=== RUN   Test_BinarySearchTree
    bstree_test.go:35: r=nil,s=0,v=0
    bstree_test.go:42: r=5,s=1,v=1 5(n,n)
    bstree_test.go:60: r=nil,s=0,v=2
    bstree_test.go:72: r=5,s=9,v=11 5(4,8),4(2,n),8(7,9),2(1,3),7(6,n),9(n,n),1(n,n),3(n,n),6(n,n)
    bstree_test.go:86: ok=true, v=5
    bstree_test.go:86: ok=true, v=4
    bstree_test.go:86: ok=true, v=8
    bstree_test.go:86: ok=true, v=2
    bstree_test.go:86: ok=true, v=7
    bstree_test.go:86: ok=true, v=9
    bstree_test.go:86: ok=true, v=1
    bstree_test.go:86: ok=true, v=3
    bstree_test.go:86: ok=true, v=6
    bstree_test.go:108: i=1 v=1 tree=r=5,s=8,v=12 5(4,8),4(2,n),8(7,9),2(n,3),7(6,n),9(n,n),3(n,n),6(n,n)
    bstree_test.go:108: i=2 v=2 tree=r=5,s=7,v=13 5(4,8),4(3,n),8(7,9),3(n,n),7(6,n),9(n,n),6(n,n)
    bstree_test.go:108: i=3 v=3 tree=r=5,s=6,v=14 5(4,8),4(n,n),8(7,9),7(6,n),9(n,n),6(n,n)
    bstree_test.go:108: i=4 v=4 tree=r=5,s=5,v=15 5(n,8),8(7,9),7(6,n),9(n,n),6(n,n)
    bstree_test.go:108: i=5 v=5 tree=r=8,s=4,v=16 8(7,9),7(6,n),9(n,n),6(n,n)
    bstree_test.go:108: i=6 v=6 tree=r=8,s=3,v=17 8(7,9),7(n,n),9(n,n)
    bstree_test.go:108: i=7 v=7 tree=r=8,s=2,v=18 8(n,9),9(n,n)
    bstree_test.go:108: i=8 v=8 tree=r=9,s=1,v=19 9(n,n)
    bstree_test.go:108: i=9 v=9 tree=r=nil,s=0,v=20
    bstree_test.go:117: r=5,s=9,v=29 5(4,8),4(2,n),8(7,9),2(1,3),7(6,n),9(n,n),1(n,n),3(n,n),6(n,n)
    bstree_test.go:120: i=1 v=9 tree=r=5,s=8,v=30 5(4,8),4(2,n),8(7,n),2(1,3),7(6,n),1(n,n),3(n,n),6(n,n)
    bstree_test.go:120: i=2 v=8 tree=r=5,s=7,v=31 5(4,7),4(2,n),7(6,n),2(1,3),6(n,n),1(n,n),3(n,n)
    bstree_test.go:120: i=3 v=7 tree=r=5,s=6,v=32 5(4,6),4(2,n),6(n,n),2(1,3),1(n,n),3(n,n)
    bstree_test.go:120: i=4 v=6 tree=r=5,s=5,v=33 5(4,n),4(2,n),2(1,3),1(n,n),3(n,n)
    bstree_test.go:120: i=5 v=5 tree=r=4,s=4,v=34 4(2,n),2(1,3),1(n,n),3(n,n)
    bstree_test.go:120: i=6 v=4 tree=r=2,s=3,v=35 2(1,3),1(n,n),3(n,n)
    bstree_test.go:120: i=7 v=3 tree=r=2,s=2,v=36 2(1,n),1(n,n)
    bstree_test.go:120: i=8 v=2 tree=r=1,s=1,v=37 1(n,n)
    bstree_test.go:120: i=9 v=1 tree=r=nil,s=0,v=38
    bstree_test.go:130: r=2,s=3,v=41 2(1,3),1(n,n),3(n,n)
    bstree_test.go:135: r=nil,s=0,v=42
    bstree_test.go:151: r=nil,s=0,v=18146
--- PASS: Test_BinarySearchTree (0.01s)
PASS
ok      command-line-arguments  0.012s

IComparator.go

定义值比较函数. 值比较函数可以返回小于, 等于, 大于三种情况

package binary_search_tree

type IComparator interface {
    Compare(a interface{}, b interface{}) CompareResult
}

type CompareResult int
const LESS CompareResult = -1
const EQUAL CompareResult = 0
const GREATER CompareResult = 1

IBinarySearchTree.go

定义二叉查找树的接口, 增删改查都要.

package binary_search_tree

type IBinarySearchTree interface {
    Size() int
    IsEmpty() bool
    IsNotEmpty() bool

    Push(value interface{})
    Min() (bool, interface{})
    Max() (bool, interface{})
    Has(value interface{}) bool
    PopMin() (bool, interface{})
    PopMax() (bool, interface{})
    Remove(value interface{}) bool
    Clear()

    Iterator() IIterator
    String() string
}

IIterator.go

定义二叉查找树的遍历接口.

package binary_search_tree

type IIterator interface {
    More() bool
    Next() (bool,interface{})
}

tComparator.go

值比较函数的包装器, 实现IComparator接口. 具体比较函数由外部传入.

package binary_search_tree

import "errors"

type FnCompare func(a interface{}, b interface{}) CompareResult

type tComparator struct {
    fnCompare FnCompare
}

func NewComparator(fn FnCompare) IComparator {
    if fn == nil {
        panic(gNullArgumentError)
    }

    return &tComparator{
        fnCompare: fn,
    }
}

func (me *tComparator) Compare(a interface{}, b interface{}) CompareResult {
    if a == nil || b == nil {
        panic(gNullArgumentError)
    }
    return me.fnCompare(a, b)
}

var gNullArgumentError = errors.New("null argument error")

tBinarySearchTree.go

二叉查找树的示例, 实现IBinarySearchTree接口.

package binary_search_tree

import (
    "fmt"
    "strings"
)

type tBinarySearchTree struct {
    comparator IComparator
    root       *tBinaryNode
    size       int
    version    uint64
}

func NewBinarySearchTree(comparator IComparator) IBinarySearchTree {
    return &tBinarySearchTree{
        comparator: comparator,
        root:       nil,
        size:       0,
        version:    0,
    }
}

func (me *tBinarySearchTree) Size() int {
    return me.size
}

func (me *tBinarySearchTree) IsEmpty() bool {
    return me.size <= 0
}

func (me *tBinarySearchTree) IsNotEmpty() bool {
    return !me.IsEmpty()
}

func (me *tBinarySearchTree) Push(value interface{}) {
    if me.IsEmpty() {
        me.root = me.append(value)
        return
    }

    for node := me.root; node != nil; {
        r := me.comparator.Compare(value, node.value)

        switch r {
        case EQUAL:
            return

        case LESS:
            if node.left == nil {
                node.LChild(me.append(value))
                return
            } else {
                node = node.left
            }
            break

        case GREATER:
            if node.right == nil {
                node.RChild(me.append(value))
                return
            } else {
                node = node.right
            }
            break
        }
    }
}

func (me *tBinarySearchTree) append(value interface{}) *tBinaryNode {
    me.size++
    me.version++
    return newBinaryNode(value)
}

func (me *tBinarySearchTree) Min() (bool, interface{}) {
    ok, node := me.getMinNode(me.root)
    if !ok {
        return false, nil
    }
    return true, node.value
}

func (me *tBinarySearchTree) Max() (bool, interface{}) {
    ok, node := me.getMaxNode(me.root)
    if !ok {
        return false, nil
    }
    return true, node.value
}


func (me *tBinarySearchTree) Has(value interface{}) bool {
    ok, _ := me.find(value)
    return ok
}

func (me *tBinarySearchTree) find(value interface{}) (ok bool, node *tBinaryNode) {
    if me.IsEmpty() {
        return false, nil
    }

    for node = me.root; node != nil; {
        r := me.comparator.Compare(value, node.value)

        switch r {
        case EQUAL:
            return true, node

        case LESS:
            if node.left == nil {
                return false, nil
            } else {
                node = node.left
            }
            break

        case GREATER:
            if node.right == nil {
                return false, nil
            } else {
                node = node.right
            }
            break
        }
    }

    return false, nil
}


func (me *tBinarySearchTree) getMinNode(from *tBinaryNode) (ok bool, left *tBinaryNode) {
    if from == nil {
        return false, nil
    }

    left = from.left
    if left == nil {
        return true, from
    }

    for {
        if left.left == nil {
            return true, left
        }

        left = left.left
    }
}

func (me *tBinarySearchTree) getMaxNode(from *tBinaryNode) (ok bool, right *tBinaryNode) {
    if from == nil {
        return false, nil
    }

    right = from.right
    if right == nil {
        return true, from
    }

    for {
        if right.right == nil {
            return true, right
        }

        right = right.right
    }
}

func (me *tBinarySearchTree) PopMin() (bool, interface{}) {
    ok, node := me.getMinNode(me.root)
    if !ok {
        return false, nil
    }
    value := node.value

    me.removeNode(node)

    me.size--
    me.version++

    return true, value
}


func (me *tBinarySearchTree) removeNode(node *tBinaryNode) {
    //var pv interface{} = "nil"
    //if node.parent != nil {
    //    pv = node.parent.value
    //}
    //fmt.Printf("removeNode %v.%v\n", pv, node.value)

    left := node.left
    right := node.right

    if node.parent == nil {
        switch node.childrenCount() {
        case 0:
            me.root = nil
            break

        case 1:
            if left != nil {
                me.root = left
                left.SetParent(nil)
            } else if right != nil {
                me.root = right
                right.SetParent(nil)
            }

        default:
            _, n := me.getMaxNode(left)
            if n == left {
                me.root = n
                n.SetParent(nil)
                n.RChild(right)

            } else {
                me.root.value = n.value
                me.removeNode(n)
            }
        }

    } else {
        parent := node.parent

        switch node.childrenCount() {
        case 0:
            if parent.left == node {
                parent.left = nil
            } else {
                parent.right = nil
            }
            break

        case 1:
            if left != nil {
                if parent.left == node {
                    parent.LChild(left)
                } else {
                    parent.RChild(left)
                }
            } else if right != nil {
                if parent.left == node {
                    parent.LChild(right)
                } else {
                    parent.RChild(right)
                }
            }

        default:
            _, n := me.getMaxNode(left)
            if n == left {
                if parent.left == node {
                    parent.LChild(n)
                } else {
                    parent.RChild(n)
                }
                n.RChild(right)

            } else {
                node.value = n.value
                me.removeNode(n)
            }
        }
    }
}

func (me *tBinarySearchTree) PopMax() (bool, interface{}) {
    ok, node := me.getMaxNode(me.root)
    if !ok {
        return false, nil
    }

    me.removeNode(node)

    me.size--
    me.version++

    return true, node.value
}

func (me *tBinarySearchTree) Remove(value interface{}) bool {
    ok, node := me.find(value)
    if !ok {
        return false
    }

    me.removeNode(node)

    me.size--
    me.version++

    return true
}

func (me *tBinarySearchTree) Clear() {
    if me.IsEmpty() {
        return
    }

    me.root = nil
    me.size = 0
    me.version++
}

func (me *tBinarySearchTree) Iterator() IIterator {
    return newBSTreeIterator(me)
}

func (me *tBinarySearchTree) String() string {
    queue := newVisitQeuue()
    queue.push(me.root)

    items := make([]string, 0)
    for ;queue.more(); {
        ok, node := queue.poll()
        if ok {
            queue.push(node.left)
            queue.push(node.right)

            var lv interface{} = "n"
            if node.left != nil {
                lv = node.left.value
            }
            var rv interface{} = "n"
            if node.right != nil {
                rv = node.right.value
            }
            items = append(items, fmt.Sprintf("%v(%v,%v)", node.value, lv,rv))

        } else {
            break
        }
    }

    r := "nil"
    if me.root != nil {
        r = fmt.Sprintf("%v", me.root.value)
    }

    if len(items) > 0 {
        return fmt.Sprintf("r=%v,s=%v,v=%v %s", r, me.size, me.version, strings.Join(items, ","))
    } else {
        return fmt.Sprintf("r=%v,s=%v,v=%v", r, me.size, me.version)
    }
}

tBinaryNode.go

二叉查找树节点

package binary_search_tree

type tBinaryNode struct {
    value interface{}
    parent *tBinaryNode
    left *tBinaryNode
    right *tBinaryNode
}

func newBinaryNode(value interface{}) *tBinaryNode {
    return &tBinaryNode{
        value: value,
        left: nil,
        right: nil,
    }
}

func (me *tBinaryNode) childrenCount() int {
    left := me.left
    right := me.right

    if left == nil && right == nil {
        return 0
    }

    if left == nil || right == nil {
        return 1
    }

    return 2
}

func (me *tBinaryNode) LChild(node *tBinaryNode) {
    me.left = node
    if node != nil {
        node.SetParent(me)
    }
}

func (me *tBinaryNode) SetParent(parent *tBinaryNode) {
    if me.parent == parent {
        return
    }

    if me.parent != nil {
        if me.parent.left == me {
            me.parent.left = nil
        } else if me.parent.right == me {
            me.parent.right = nil
        }
    }

    me.parent = parent
}

func (me *tBinaryNode) RChild(node *tBinaryNode) {
    me.right = node
    if node != nil {
        node.SetParent(me)
    }
}

tBSTreeIterator.go

二叉查找树的遍历迭代器, 内部使用广度优先搜索+候选队列.

package binary_search_tree

type tBSTreeIterator struct {
    tree *tBinarySearchTree
    queue *tVisitQueue
    version uint64
}


type tVisitQueue struct {
    head *tQueuedNode
    tail *tQueuedNode
}

type tQueuedNode struct {
    node *tBinaryNode
    next *tQueuedNode
}

func newQueuedNode(node *tBinaryNode) *tQueuedNode {
    return &tQueuedNode{
        node: node,
        next: nil,
    }
}

func newVisitQeuue() *tVisitQueue {
    return &tVisitQueue{
        nil,
        nil,
    }
}

func (me *tVisitQueue) push(node *tBinaryNode) {
    if node == nil {
        return
    }

    qn := newQueuedNode(node)
    if me.head == nil {
        me.head = qn
        me.tail = qn
    } else {
        me.tail.next = qn
        me.tail = qn
    }
}

func (me *tVisitQueue) poll() (bool, *tBinaryNode) {
    if me.head == nil {
        return false, nil

    } else {
        it := me.head
        if it == me.tail {
            me.tail = nil
        }
        me.head = me.head.next
        return true, it.node
    }
}

func (me *tVisitQueue) more() bool {
    return me.head != nil
}

func newBSTreeIterator(tree *tBinarySearchTree) IIterator {
    it := &tBSTreeIterator{
        tree: tree,
        queue: newVisitQeuue(),
        version: tree.version,
    }
    it.queue.push(tree.root)
    return it
}

func (me *tBSTreeIterator) More() bool {
    if me.version != me.tree.version {
        return false
    } else {
        return me.queue.more()
    }
}

func (me *tBSTreeIterator) Next() (bool, interface{}) {
    if !me.More() {
        return false, nil
    }

    ok, node := me.queue.poll()
    if !ok {
        return false, nil
    }

    me.queue.push(node.left)
    me.queue.push(node.right)
    return true, node.value
}

(end)

有疑问加站长微信联系（非本文作者）