二分木(Binary Search Tree)の実装 その1

概要

シンプルな二分木をGolangで実装していきます。 二分木の動作は以下のURLがビジュアル化されていてわかりやすいです。

Binary Search Tree Visualization

まずは二分木のベースとして値の挿入をして木を作成し、 見つけたい値が木に含まれているかを判定する関数を作成します。

実装

木とノードの定義

type Node struct {
    Value int  // ノードの持つ値
    Left  *Node // 左側ノードへのポインタ
    Right *Node // 右側ノードへのポインタ
}

// 二分木として持つのはルートノードのみ
// ノードにアクセスするのはすべてルートからたどる
type Tree struct {
    Root *Node
}

代入処理

// TreeにNodeを加える
func (t *Tree) Insert(insertValue int) {
    if t.Root == nil {
        // 初回のNodeだった場合はTreeのRootとしてNodeを設定する
        // t.root = CreateNode(insertValue)
        t.Root = &Node{Value: insertValue}
    } else {
        parentNode := t.Root
        for {
            if insertValue >= parentNode.Value {
                // 右のノードを確認
                if parentNode.Right != nil {
                    // ノードが存在するならば移動
                    parentNode = parentNode.Right
                } else {
                    // ノードがまだ無ければノードを配置する
                    parentNode.Right = &Node{Value: insertValue}
                    break
                }
            } else {
                // 左のノードを確認して同様に代入OR移動
                if parentNode.Left != nil {
                    parentNode = parentNode.Left
                } else {
                    parentNode.Left = &Node{Value: insertValue}
                    break
                }
            }
        }
    }
}

テストの作成

値を代入して正しい位置にノードがあるかを判定するテストコードを作成し、正しく処理ができているか確認します。

ここでは 4, 2, 7, 6 の順に代入して目的の木ができているかを確認します。

f:id:itib:20211017145547p:plain
4, 2, 7, 6の順で代入した木

func TestCreateTree(t *testing.T) {
    var tree binary_search_tree.Tree
    tree.Insert(4)
    tree.Insert(2)
    tree.Insert(7)
    tree.Insert(6)

    if tree.Root.Value != 4 {
        t.Error("Failed")
    }
    if tree.Root.Left.Value != 2 {
        t.Error("Failed")
    }
    if tree.Root.Right.Value != 7 {
        t.Error("Failed")
    }
    if tree.Root.Right.Left.Value != 6 {
        t.Error("Failed")
    }
    if tree.Root.Right.Right != nil {
        t.Error("Failed")
    }
}

ディレクトリ構成

project
├ binary_search_tree
│ ├ binary_search_tree.go
│ └ binary_search_tree_test.go
└ go.mod

binary_search_tree.go

package binary_search_tree

type Node struct {
    Value int
    Left  *Node
    Right *Node
}

// NodeたちにアクセスするのはすべてRootから
// BinarySearchTreeとして知っているのはRootの場所のみ
type Tree struct {
    Root *Node
}

// TreeにNodeを加える
func (t *Tree) Insert(insertValue int) {
    if t.Root == nil {
        // 初回のNodeだった場合はTreeのRootとしてNodeを設定する
        // t.root = CreateNode(insertValue)
        t.Root = &Node{Value: insertValue}
    } else {
        parentNode := t.Root
        for {
            if insertValue >= parentNode.Value {
                // 右のノードを確認
                if parentNode.Right != nil {
                    // ノードが存在するならば移動
                    parentNode = parentNode.Right
                } else {
                    // ノードがまだ無ければノードを配置する
                    parentNode.Right = &Node{Value: insertValue}
                    break
                }
            } else {
                // 左のノードを確認して同様に代入OR移動
                if parentNode.Left != nil {
                    parentNode = parentNode.Left
                } else {
                    parentNode.Left = &Node{Value: insertValue}
                    break
                }
            }
        }
    }
}

binary_search_tree_test.go

package binary_search_tree_test

import (
    "testing"

    "github.com/itiB/project/binary_search_tree"
)

func TestCreateTree(t *testing.T) {
    var tree binary_search_tree.Tree
    tree.Insert(4)
    tree.Insert(2)
    tree.Insert(7)
    tree.Insert(6)

    if tree.Root.Value != 4 {
        t.Error("Failed")
    }
    if tree.Root.Left.Value != 2 {
        t.Error("Failed")
    }
    if tree.Root.Right.Value != 7 {
        t.Error("Failed")
    }
    if tree.Root.Right.Left.Value != 6 {
        t.Error("Failed")
    }
    if tree.Root.Right.Right != nil {
        t.Error("Failed")
    }
}

go.mod

module github.com/itiB/project

go 1.16

テスト結果正常に木が作成できていることがわかりました。

❯ go test ./binary_search_tree
ok      github.com/itiB/project/binary_search_tree   0.004s

検索したい値が木に含まれているか判定

検索対象の値が木に含まれているか調べるには代入する場所を探すのと同様に木を探索していって一致する値が含まれているかを調べることで実装できます。

func (tree *Tree) Contains(findValue int) bool {
    // そもそも木がなければ含まれていない
    if tree.Root == nil {
        return false
    }

    cursor := tree.Root // 調べているノードを指す
    for {
        // ノードの値が検索対象の値なら発見
        if cursor.Value == findValue {
            return true
        } else {
            // ノードを移動する処理
            // もし進むべき方向にノードがないならば値は含まれていないとしてfalse
            if findValue >= cursor.Value {
                if cursor.Right != nil {
                    cursor = cursor.Right
                } else {
                    return false
                }
            } else {
                if cursor.Left != nil {
                    cursor = cursor.Left
                } else {
                    return false
                }
            }
        }
    }
}

テストの作成

木を作成し、代入した値が含まれているかのテストを記述します。

func TestFindTree(t *testing.T) {
    var tree binary_search_tree.Tree
    tree.Insert(4)
    tree.Insert(2)
    tree.Insert(7)
    tree.Insert(6)
    tree.Insert(-5)

    if !tree.Contains(4) {
        t.Errorf("Failed")
    }
    if !tree.Contains(6) {
        t.Errorf("Failed")
    }
    if !tree.Contains(-5) {
        t.Errorf("Failed")
    }
    if tree.Contains(5) {
        t.Errorf("Failed")
    }
    if tree.Contains(500000000000) {
        t.Errorf("Failed")
    }
}
❯ go test ./binary_search_tree
ok      github.com/itiB/project/binary_search_tree   0.004s

まとめ

二分木のベースとなる処理を作成することができました。 次は削除処理を作成します。