二分探索木

今日は、練習のため二分探索木の写経をする。

cf. 二分探索木 – M.Hiroi’s Home Page お気楽 Go 言語プログラミング入門

package main

import "fmt"

type Item interface {
    Eq(Item) bool
    Less(Item) bool
}

type Node struct {
    item Item
    left, right *Node
}

func newNode(x Item) *Node {
    p := new(Node)
    p.item = x
    return p
}

func searchNode(node *Node, x Item) bool {
    for node != nil {
        switch {
            case x.Eq(node.item): return true
            case x.Less(node.item): node = node.left
            default: node = node.right
        }
    }
    return false
}

func insertNode(node *Node, x Item) *Node {
    switch {
        case node == nil: return newNode(x)
        case x.Eq(node.item): return node
        case x.Less(node.item):
            node.left = insertNode(node.left, x)
        default:
            node.right = insertNode(node.right, x)
    }
    return node
}

func searchMin(node *Node) Item {
    if node.left == nil {
        return node.item
    }
    return searchMin(node.left)
}

func deleteMin(node *Node) *Node {
    if node.left == nil {
        return node.right
    }
    node.left = deleteMin(node.left)
    return node
}

func deleteNode(node *Node, x Item) *Node {
    if node != nil {
        if x.Eq(node.item) {
            if node.left == nil {
                return node.right
            } else if node.right == nil {
                return node.left
            } else {
                node.item = searchMin(node.right)
                node.right = deleteMin(node.right)
            }
        } else if x.Less(node.item) {
            node.left = deleteNode(node.left, x)
        } else {
            node.right = deleteNode(node.right, x)
        }
    }
    return node
}

func foreachNode(f func(Item), node *Node) {
    if node != nil {
        foreachNode(f, node.left)
        f(node.item)
        foreachNode(f, node.right)
    }
}

type Tree struct {
    root *Node
}

func newTree() *Tree {
    return new(Tree)
}

func (t *Tree) searchTree(x Item) bool {
    return searchNode(t.root, x)
}

func (t *Tree) insertTree(x Item) {
    t.root = insertNode(t.root, x)
}

func (t *Tree) deleteTree(x Item) {
    t.root = deleteNode(t.root, x)
}

func (t *Tree) foreachTree(f func(Item)) {
    foreachNode(f, t.root)
}

func (t *Tree) printTree() {
    t.foreachTree(func(x Item) { fmt.Print(x, " ") })
    fmt.Println("")
}

type Int int

func (n Int) Eq(m Item) bool {
    return n == m.(Int)
}

func (n Int) Less(m Item) bool {
    return n < m.(Int)
}

func main() {
    a := newTree()
    for _, v := range []int { 5,6,4,7,3,8,2,9,1,0 } {
        a.insertTree(Int(v))
    }
    a.printTree()
    for i := 0; i < 10; i++ {
        fmt.Println(a.searchTree(Int(i)))
    }
    a.printTree()
    for i := 0; i < 10; i++ {
        a.deleteTree(Int(i)) a.printTree()
    }
}

探索木自体は Item 型(インターフェイス)を対象としている。このため intInt という別名をつけて EqLess の2つのメソッドを定義している。わざわざ別名をつけているのは、通常の(構造体でない)型にはメソッドを定義できないため。

^o^ > go run binarytree.go
0 1 2 3 4 5 6 7 8 9
true
true
true
true
true
true
true
true
true
true
0 1 2 3 4 5 6 7 8 9
1 2 3 4 5 6 7 8 9
2 3 4 5 6 7 8 9
3 4 5 6 7 8 9
4 5 6 7 8 9
5 6 7 8 9
6 7 8 9
7 8 9
8 9
9