Goで実装するBtree[挿入・探索編]

package tree

// Item - ノードに属するアイテム
type Item struct {
    Value  int
    Object interface{}
}

// Items - アイテムのリスト
type Items []*Item

func (i *Item) Less(item *Item) bool {
    return i.Value < item.Value
}

// Node - ノード。属するアイテムのリストと、子になるノード群を持っている
type Node struct {
    Items  Items
    Branch Branch
}

// Branch - あるノードに属するノード群
type Branch []*Node

// Btree - B-treeの本体。ルートになるノードへのポインタを持つ
type Btree struct {
    Root *Node
}

// NewBTree - B-treeを作成する。mはB-treeのオーダー
func NewBTree(m int) (*Btree, error) {
    if m < 2 {
        return nil, fmt.Errorf("m must be greater than 3, can't be %d", m)
    }

    return &Btree{
        m:    m,
        Root: nil,
    }, nil
}

// maxItemNum - アイテムの閾値を返す
func (b *Btree) maxItemNum() int {
    return b.m*2 - 1
}

tree := btree.NewBtree(2)

1. ルートノードのアイテムリストから対象なアイテムを探索
2. 1で見つからなければ子ノードを探索
3. 子ノードのアイテムリストを探索し、なければその子ノードを(以下略)

// findItem - Itemsの中から検索対象のItemを検索する。なかった場合は、もっとも近く、小さいItemのindexを返す
func (i *Items) findItem(item *Item) (index int, exist bool) {
    index = sort.Search(len(*i), func(in int) bool {
        return item.Less((*i)[in])
    })

    // あるItemと等価だった場合は、一つ前のindexで返す
    if index > 0 && !((*i)[index-1].Less(item)) {
        return index - 1, true
    }

    return index, false
}

// findItem - ノードの要素、Branchの中に指定されたItemがあるかをチェック
func (n *Node) findItem(item *Item) (*Item, bool) {
    index, ok := n.Items.findItem(item)
    if ok {
        return n.Items[index], true
    }

    if len((*n).Branch) > 0 {
        return n.Branch[index].findItem(item)
    }

    return nil, false
}

func (b *Btree) Find(item *Item) (*Item, bool) {
    return b.Root.findItem(item)
}

func main() {
    btree := tree.Btree{
        Root: &tree.Node{
            Items: tree.Items{
                {Value: 10, Object: "田中"},
            },
            Branch: tree.Branch{
                {Items: tree.Items{
                    {Value: 4, Object: "佐藤"},
                    {Value: 6, Object: "鈴木"},
                }},
                {Items: tree.Items{
                    {Value: 15, Object: "山田"},
                    {Value: 16, Object: "斎藤"},
                }},
            },
        },
    }
    if item, ok := btree.Find(&tree.Item{Value: 15}); ok {
        fmt.Println(item.Object)
    }
}

// output: 山田

1.ルートノードに余裕があればルートノードのアイテムリストに追加
2. 閾値(m)を超える場合は、自分自身を分割し、自身の子ノードに追加する
3. 子ノードが閾値を超える場合は、子ノードを分割し、アイテムリストに真ん中になるアイテムを追加し、子ノードに分割したノードを追加する。
4. その結果2の条件に当てはまれば、さらに自分自身を分割する。

func (i *Items) insertAt(item *Item, index int) {
    // 枠を一個増やす
    *i = append(*i, nil)

    // 最後のindex以外だった場合、新しいItemを入れるスペースを作る
    if index < len(*i) {
        copy((*i)[index+1:], (*i)[index:])
    }

    (*i)[index] = item
}

1. すでにアイテムがあれば置き換える
2. 子ノードがない場合には、アイテムリストに追加する。
3. アイテムリストが閾値を超えた場合は真ん中で分割する。

func (n *Node) shouldSplit(i, limit int) bool {
    if len(n.Branch[i].Items) <= limit {
        return false
    }

    // limitを超えて分割すべきだった場合、分割してBranchに登録
    center, newNode := n.Branch[i].split(limit / 2)
    n.Items.insertAt(center, i)
    n.Branch.InsertAt(i+1, newNode)

    return true
}

func (n *Node) split(index int) (*Item, *Node) {
    // ちょうど分割の真ん中にあるItem
    item := n.Items[index]

    newNode := new(Node)
    newNode.Items = append(newNode.Items, n.Items[index+1:]...)

    n.Items.extract(index)
    if len(n.Branch) > 0 {
        newNode.Branch = append(newNode.Branch, n.Branch[index+1:]...)
        n.Branch.extract(index + 1)
    }

    return item, newNode
}

// extract - 指定されたindexまでのItemのみを残す
func (i *Items) extract(index int) {
    *i = (*i)[:index]
}

// extract - 指定された子ノードのみを残す
func (b *Branch) extract(index int) {
    *b = (*b)[:index]

    // index以降を初期化
    for i := index; i < len(*b); i++ {
        (*b)[i] = nil
    }
}

//insert - nodeにItemを挿入する。limitの数を超えていた場合、新しくNodeを作成する
//Itemがもしあれば、Itemを入れ替え、入れ替えるまえのItemを返却する
func (n *Node) insert(item *Item, limit int) *Item {
    i, ok := n.Items.findItem(item)
    if ok {
        res := n.Items[i]
        n.Items[i] = item
        return res
    }

    if len(n.Branch) == 0 {
        n.Items.insertAt(item, i)
        return nil
    }

    // 子ノードが閾値を超える場合は、分割し、アイテムリストに追加する
    if n.shouldSplit(i, limit) {
        inTree := n.Items[i]
        switch {
        case item.Less(inTree):
            // no change, we want first split node
        case inTree.Less(item):
            i++ // we want second split node
        default:
            out := n.Items[i]
            n.Items[i] = item
            return out
        }
    }

    return n.Branch[i].insert(item, limit)
}

func (b *Btree) InsertOrUpdateItem(item *Item) (*Item, error) {
    if item == nil {
        return nil, fmt.Errorf("nil item is unacceptable")
    }

    // RootのNodeがまだなければ作成
    if b.Root == nil {
        b.Root = &Node{}
        b.Root.Items = append(b.Root.Items, item)
        return nil, nil
    } else {
        // Rootに収まらなくなったら分割する
        if len(b.Root.Items) >= b.maxItemNum() {
            item2, second := b.Root.split(b.maxItemNum() / 2)
            old := b.Root

            b.Root = new(Node)
            b.Root.Items = append(b.Root.Items, item2)
            b.Root.Branch = append(b.Root.Branch, old, second)
        }
    }

    res := b.Root.insert(item, b.maxItemNum())
    return res, nil
}