go で赤黒木

red-black-tree.go

package main

import (
	"fmt"
	"strings"
)

type Ord int

const (
	LT Ord = iota
	EQ Ord = iota
	GT Ord = iota
)

type Color interface {
	balance(left Tree, value int, right Tree) Tree
	ToString(value int) string
	IsRed() bool
}

func main() {
	tree := empty()
	numbers := []int { 6,1,0,5,4,7,2,9,3 }
	for _, i := range numbers {
		tree = tree.Insert(i)
		fmt.Println("====", i, "====")
		fmt.Println(tree.ToString(0))
	}
	fmt.Println(tree.ToString(0))
}

func empty() Tree {
	return LEAF
}

type Tree interface {
	ToBlack() Tree
	Insert(newValue int) Tree
	insertPlain(newValue int) Tree
	isLeaf() bool
	ToString(depth int) string
}

type Leaf struct {
}

var LEAF = &Leaf{}

func (l *Leaf) isLeaf() bool {
	return true
}


func (l *Leaf) ToBlack() Tree {
	return l
}

func (l *Leaf) Insert(newValue int) Tree {
	return l.insertPlain(newValue).ToBlack()
}

func (l *Leaf) insertPlain(newValue int) Tree {
	return &Node{Color:red(), Value: newValue, Left: l, Right: l}
}

func (l *Leaf) ToString(depth int) string {
	space := strings.Repeat("  ", depth)
	return fmt.Sprintf("%s-empty-", space)
}

type Node struct {
	Color Color
	Value int

	Left  Tree
	Right Tree
}

func (node *Node) isRed() bool {
	return node.Color.IsRed()
}

func (node *Node) isLeaf() bool {
	return false
}

func (node *Node) ToBlack() Tree {
	return &Node{Color: black(), Value: node.Value, Left: node.Left, Right: node.Right}
}

func (node *Node) ToString(depth int) string {
	space := strings.Repeat("  ", depth)
	return fmt.Sprintf("%s\n%s%s\n%s", node.Left.ToString(depth+1), space, node.Color.ToString(node.Value), node.Right.ToString(depth+1))
}

func (node *Node) Insert(newValue int) Tree {
	return node.insertPlain(newValue).ToBlack()
}

func (node *Node) insertPlain(newValue int) Tree {
	if newValue < node.Value {
		newTree := node.Color.balance(node.Left.insertPlain(newValue), node.Value, node.Right)
		return newTree
	} else if node.Value < newValue {
		newTree := node.Color.balance(node.Left, node.Value, node.Right.insertPlain(newValue))
		return newTree
	} else {
		return node
	}
}

type Black struct {
}

var BLACK = &Black{}

func (b *Black) ToString(value int) string {
	return fmt.Sprintf("B(%d)", value)
}

func (*Black) IsRed() bool {
	return false
}

func black() Color {
	return BLACK
}

func (b *Black) balance(left Tree, value int, right Tree) Tree {
	if left.isLeaf() && right.isLeaf() {
		return &Node{Color:b, Value: value, Left: left, Right: right}
	}
	if !left.isLeaf() {
		leftNode := left.(*Node)
		if !leftNode.Left.isLeaf() {
			leftSubLeft := leftNode.Left.(*Node)
			if leftNode.isRed() && leftSubLeft.isRed() {
				return &Node{
					Color: red(),
					Value: leftNode.Value,
					Left:  leftSubLeft.ToBlack(),
					Right: &Node{Color:black(), Value:value, Left: leftNode.Right, Right:right},
				}
			}
		} else if !leftNode.Right.isLeaf() {
			leftSubRight := leftNode.Right.(*Node)
			if leftNode.isRed() && leftSubRight.isRed() {
				return &Node{
					Color:red(),
					Value:leftSubRight.Value,
					Left:&Node{Color:black(), Value:leftNode.Value, Left:leftNode.Left, Right:leftSubRight.Left},
					Right:&Node{Color:black(), Value:value, Left:leftSubRight.Right, Right:right},
				}
			}
		} else if right.isLeaf() {
			return &Node{Color:b, Value: value, Left: left, Right: right}
		}
	}
	rightNode := right.(*Node)
	if !rightNode.Right.isLeaf() {
		rightSubRight := rightNode.Right.(*Node)
		if rightNode.isRed() && rightSubRight.isRed() {
			return &Node{
				Color:red(),
				Value:rightNode.Value,
				Left:&Node{Color:black(), Value: value, Left: left, Right:rightNode.Left},
				Right: rightSubRight.ToBlack(),
			}
		}
	} else if !rightNode.Left.isLeaf() {
		rightSubLeft := rightNode.Left.(*Node)
		if rightNode.isRed() && rightSubLeft.isRed() {
			return &Node{
				Color:red(),
				Value:rightSubLeft.Value,
				Left:&Node{Color:black(), Value:value, Left:left, Right:rightSubLeft.Left},
				Right:&Node{Color:black(), Value:rightNode.Value, Left:rightSubLeft.Right, Right:rightNode.Right},
			}
		}
	}
	return &Node{Color:b, Value: value, Left: left, Right: right}
}

type Red struct {
}

var RED = &Red{}

func (r *Red) ToString(value int) string {
	return fmt.Sprintf("R(%d)", value)
}

func (*Red) IsRed() bool {
	return true
}

func red() Color {
	return RED
}

func (r *Red) balance(left Tree, value int, right Tree) Tree {
	return &Node{Color: red(), Value: value, Left: left, Right: right}
}