mirror of https://github.com/emirpasic/gods
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
554 lines
12 KiB
Go
554 lines
12 KiB
Go
// Copyright (c) 2015, Emir Pasic. All rights reserved.
|
|
// Use of this source code is governed by a BSD-style
|
|
// license that can be found in the LICENSE file.
|
|
|
|
// Package redblacktree implements a red-black tree.
|
|
//
|
|
// Used by TreeSet and TreeMap.
|
|
//
|
|
// Structure is not thread safe.
|
|
//
|
|
// References: http://en.wikipedia.org/wiki/Red%E2%80%93black_tree
|
|
package redblacktree
|
|
|
|
import (
|
|
"fmt"
|
|
|
|
"github.com/emirpasic/gods/utils"
|
|
"github.com/spewspews/gods/trees"
|
|
)
|
|
|
|
func assertTreeImplementation() {
|
|
var _ trees.Tree = (*Tree)(nil)
|
|
}
|
|
|
|
type color bool
|
|
|
|
const (
|
|
black, red color = true, false
|
|
)
|
|
|
|
// Tree holds elements of the red-black tree
|
|
type Tree struct {
|
|
Root *Node
|
|
size int
|
|
comparator utils.Comparator
|
|
}
|
|
|
|
// Node is a single element within the tree
|
|
type Node struct {
|
|
Key interface{}
|
|
Value interface{}
|
|
color color
|
|
Left *Node
|
|
Right *Node
|
|
Parent *Node
|
|
}
|
|
|
|
// NewWith instantiates a red-black tree with the custom comparator.
|
|
func NewWith(comparator utils.Comparator) *Tree {
|
|
return &Tree{comparator: comparator}
|
|
}
|
|
|
|
// NewWithIntComparator instantiates a red-black tree with the IntComparator, i.e. keys are of type int.
|
|
func NewWithIntComparator() *Tree {
|
|
return &Tree{comparator: utils.IntComparator}
|
|
}
|
|
|
|
// NewWithStringComparator instantiates a red-black tree with the StringComparator, i.e. keys are of type string.
|
|
func NewWithStringComparator() *Tree {
|
|
return &Tree{comparator: utils.StringComparator}
|
|
}
|
|
|
|
// Comparator returns the comparator function for the tree.
|
|
func (tree *Tree) Comparator() utils.Comparator {
|
|
return tree.comparator
|
|
}
|
|
|
|
// New returns a new empty tree with the same comparator.
|
|
func (tree *Tree) New() trees.Tree {
|
|
return &Tree{comparator: tree.comparator}
|
|
}
|
|
|
|
// Put inserts node into the tree.
|
|
// Key should adhere to the comparator's type assertion, otherwise method panics.
|
|
func (tree *Tree) Put(key interface{}, value interface{}) {
|
|
insertedNode := &Node{Key: key, Value: value, color: red}
|
|
if tree.Root == nil {
|
|
tree.Root = insertedNode
|
|
} else {
|
|
node := tree.Root
|
|
loop := true
|
|
for loop {
|
|
compare := tree.comparator(key, node.Key)
|
|
switch {
|
|
case compare == 0:
|
|
node.Key = key
|
|
node.Value = value
|
|
return
|
|
case compare < 0:
|
|
if node.Left == nil {
|
|
node.Left = insertedNode
|
|
loop = false
|
|
} else {
|
|
node = node.Left
|
|
}
|
|
case compare > 0:
|
|
if node.Right == nil {
|
|
node.Right = insertedNode
|
|
loop = false
|
|
} else {
|
|
node = node.Right
|
|
}
|
|
}
|
|
}
|
|
insertedNode.Parent = node
|
|
}
|
|
tree.insertCase1(insertedNode)
|
|
tree.size++
|
|
}
|
|
|
|
// Get searches the node in the tree by key and returns its value or nil if key is not found in tree.
|
|
// Second return parameter is true if key was found, otherwise false.
|
|
// Key should adhere to the comparator's type assertion, otherwise method panics.
|
|
func (tree *Tree) Get(key interface{}) (value interface{}, found bool) {
|
|
node := tree.lookup(key)
|
|
if node != nil {
|
|
return node.Value, true
|
|
}
|
|
return nil, false
|
|
}
|
|
|
|
// Remove remove the node from the tree by key.
|
|
// Key should adhere to the comparator's type assertion, otherwise method panics.
|
|
func (tree *Tree) Remove(key interface{}) {
|
|
var child *Node
|
|
node := tree.lookup(key)
|
|
if node == nil {
|
|
return
|
|
}
|
|
if node.Left != nil && node.Right != nil {
|
|
pred := node.Left.maximumNode()
|
|
node.Key = pred.Key
|
|
node.Value = pred.Value
|
|
node = pred
|
|
}
|
|
if node.Left == nil || node.Right == nil {
|
|
if node.Right == nil {
|
|
child = node.Left
|
|
} else {
|
|
child = node.Right
|
|
}
|
|
if node.color == black {
|
|
node.color = nodeColor(child)
|
|
tree.deleteCase1(node)
|
|
}
|
|
tree.replaceNode(node, child)
|
|
if node.Parent == nil && child != nil {
|
|
child.color = black
|
|
}
|
|
}
|
|
tree.size--
|
|
}
|
|
|
|
// Empty returns true if tree does not contain any nodes
|
|
func (tree *Tree) Empty() bool {
|
|
return tree.size == 0
|
|
}
|
|
|
|
// Size returns number of nodes in the tree.
|
|
func (tree *Tree) Size() int {
|
|
return tree.size
|
|
}
|
|
|
|
// Keys returns all keys in-order
|
|
func (tree *Tree) Keys() []interface{} {
|
|
keys := make([]interface{}, tree.size)
|
|
it := tree.Iterator()
|
|
for i := 0; it.Next(); i++ {
|
|
keys[i] = it.Key()
|
|
}
|
|
return keys
|
|
}
|
|
|
|
// Values returns all values in-order based on the key.
|
|
func (tree *Tree) Values() []interface{} {
|
|
values := make([]interface{}, tree.size)
|
|
it := tree.Iterator()
|
|
for i := 0; it.Next(); i++ {
|
|
values[i] = it.Value()
|
|
}
|
|
return values
|
|
}
|
|
|
|
// Left returns the left-most (min) node or nil if tree is empty.
|
|
func (tree *Tree) Left() *Node {
|
|
var parent *Node
|
|
current := tree.Root
|
|
for current != nil {
|
|
parent = current
|
|
current = current.Left
|
|
}
|
|
return parent
|
|
}
|
|
|
|
// Right returns the right-most (max) node or nil if tree is empty.
|
|
func (tree *Tree) Right() *Node {
|
|
var parent *Node
|
|
current := tree.Root
|
|
for current != nil {
|
|
parent = current
|
|
current = current.Right
|
|
}
|
|
return parent
|
|
}
|
|
|
|
// Min returns the minimum key value pair of the AVL tree
|
|
// or nils if the tree is empty.
|
|
func (tree *Tree) Min() (interface{}, interface{}) {
|
|
n := tree.Left()
|
|
if n == nil {
|
|
return nil, nil
|
|
}
|
|
return n.Key, n.Value
|
|
}
|
|
|
|
// Max returns the minimum key value pair of the AVL tree
|
|
// or nils if the tree is empty.
|
|
func (tree *Tree) Max() (interface{}, interface{}) {
|
|
n := tree.Right()
|
|
if n == nil {
|
|
return nil, nil
|
|
}
|
|
return n.Key, n.Value
|
|
}
|
|
|
|
// Floor Finds floor node of the input key, return the floor node or nil if no ceiling is found.
|
|
// Second return parameter is true if floor was found, otherwise false.
|
|
//
|
|
// Floor node is defined as the largest node that is smaller than or equal to the given node.
|
|
// A floor node may not be found, either because the tree is empty, or because
|
|
// all nodes in the tree is larger than the given node.
|
|
//
|
|
// Key should adhere to the comparator's type assertion, otherwise method panics.
|
|
func (tree *Tree) Floor(key interface{}) (floor *Node, found bool) {
|
|
found = false
|
|
node := tree.Root
|
|
for node != nil {
|
|
compare := tree.comparator(key, node.Key)
|
|
switch {
|
|
case compare == 0:
|
|
return node, true
|
|
case compare < 0:
|
|
node = node.Left
|
|
case compare > 0:
|
|
floor, found = node, true
|
|
node = node.Right
|
|
}
|
|
}
|
|
if found {
|
|
return floor, true
|
|
}
|
|
return nil, false
|
|
}
|
|
|
|
// Ceiling finds ceiling node of the input key, return the ceiling node or nil if no ceiling is found.
|
|
// Second return parameter is true if ceiling was found, otherwise false.
|
|
//
|
|
// Ceiling node is defined as the smallest node that is larger than or equal to the given node.
|
|
// A ceiling node may not be found, either because the tree is empty, or because
|
|
// all nodes in the tree is smaller than the given node.
|
|
//
|
|
// Key should adhere to the comparator's type assertion, otherwise method panics.
|
|
func (tree *Tree) Ceiling(key interface{}) (ceiling *Node, found bool) {
|
|
found = false
|
|
node := tree.Root
|
|
for node != nil {
|
|
compare := tree.comparator(key, node.Key)
|
|
switch {
|
|
case compare == 0:
|
|
return node, true
|
|
case compare < 0:
|
|
ceiling, found = node, true
|
|
node = node.Left
|
|
case compare > 0:
|
|
node = node.Right
|
|
}
|
|
}
|
|
if found {
|
|
return ceiling, true
|
|
}
|
|
return nil, false
|
|
}
|
|
|
|
// Clear removes all nodes from the tree.
|
|
func (tree *Tree) Clear() {
|
|
tree.Root = nil
|
|
tree.size = 0
|
|
}
|
|
|
|
// String returns a string representation of container
|
|
func (tree *Tree) String() string {
|
|
str := "RedBlackTree\n"
|
|
if !tree.Empty() {
|
|
output(tree.Root, "", true, &str)
|
|
}
|
|
return str
|
|
}
|
|
|
|
func (node *Node) String() string {
|
|
return fmt.Sprintf("%v", node.Key)
|
|
}
|
|
|
|
func output(node *Node, prefix string, isTail bool, str *string) {
|
|
if node.Right != nil {
|
|
newPrefix := prefix
|
|
if isTail {
|
|
newPrefix += "│ "
|
|
} else {
|
|
newPrefix += " "
|
|
}
|
|
output(node.Right, newPrefix, false, str)
|
|
}
|
|
*str += prefix
|
|
if isTail {
|
|
*str += "└── "
|
|
} else {
|
|
*str += "┌── "
|
|
}
|
|
*str += node.String() + "\n"
|
|
if node.Left != nil {
|
|
newPrefix := prefix
|
|
if isTail {
|
|
newPrefix += " "
|
|
} else {
|
|
newPrefix += "│ "
|
|
}
|
|
output(node.Left, newPrefix, true, str)
|
|
}
|
|
}
|
|
|
|
func (tree *Tree) lookup(key interface{}) *Node {
|
|
node := tree.Root
|
|
for node != nil {
|
|
compare := tree.comparator(key, node.Key)
|
|
switch {
|
|
case compare == 0:
|
|
return node
|
|
case compare < 0:
|
|
node = node.Left
|
|
case compare > 0:
|
|
node = node.Right
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (node *Node) grandparent() *Node {
|
|
if node != nil && node.Parent != nil {
|
|
return node.Parent.Parent
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (node *Node) uncle() *Node {
|
|
if node == nil || node.Parent == nil || node.Parent.Parent == nil {
|
|
return nil
|
|
}
|
|
return node.Parent.sibling()
|
|
}
|
|
|
|
func (node *Node) sibling() *Node {
|
|
if node == nil || node.Parent == nil {
|
|
return nil
|
|
}
|
|
if node == node.Parent.Left {
|
|
return node.Parent.Right
|
|
}
|
|
return node.Parent.Left
|
|
}
|
|
|
|
func (tree *Tree) rotateLeft(node *Node) {
|
|
right := node.Right
|
|
tree.replaceNode(node, right)
|
|
node.Right = right.Left
|
|
if right.Left != nil {
|
|
right.Left.Parent = node
|
|
}
|
|
right.Left = node
|
|
node.Parent = right
|
|
}
|
|
|
|
func (tree *Tree) rotateRight(node *Node) {
|
|
left := node.Left
|
|
tree.replaceNode(node, left)
|
|
node.Left = left.Right
|
|
if left.Right != nil {
|
|
left.Right.Parent = node
|
|
}
|
|
left.Right = node
|
|
node.Parent = left
|
|
}
|
|
|
|
func (tree *Tree) replaceNode(old *Node, new *Node) {
|
|
if old.Parent == nil {
|
|
tree.Root = new
|
|
} else {
|
|
if old == old.Parent.Left {
|
|
old.Parent.Left = new
|
|
} else {
|
|
old.Parent.Right = new
|
|
}
|
|
}
|
|
if new != nil {
|
|
new.Parent = old.Parent
|
|
}
|
|
}
|
|
|
|
func (tree *Tree) insertCase1(node *Node) {
|
|
if node.Parent == nil {
|
|
node.color = black
|
|
} else {
|
|
tree.insertCase2(node)
|
|
}
|
|
}
|
|
|
|
func (tree *Tree) insertCase2(node *Node) {
|
|
if nodeColor(node.Parent) == black {
|
|
return
|
|
}
|
|
tree.insertCase3(node)
|
|
}
|
|
|
|
func (tree *Tree) insertCase3(node *Node) {
|
|
uncle := node.uncle()
|
|
if nodeColor(uncle) == red {
|
|
node.Parent.color = black
|
|
uncle.color = black
|
|
node.grandparent().color = red
|
|
tree.insertCase1(node.grandparent())
|
|
} else {
|
|
tree.insertCase4(node)
|
|
}
|
|
}
|
|
|
|
func (tree *Tree) insertCase4(node *Node) {
|
|
grandparent := node.grandparent()
|
|
if node == node.Parent.Right && node.Parent == grandparent.Left {
|
|
tree.rotateLeft(node.Parent)
|
|
node = node.Left
|
|
} else if node == node.Parent.Left && node.Parent == grandparent.Right {
|
|
tree.rotateRight(node.Parent)
|
|
node = node.Right
|
|
}
|
|
tree.insertCase5(node)
|
|
}
|
|
|
|
func (tree *Tree) insertCase5(node *Node) {
|
|
node.Parent.color = black
|
|
grandparent := node.grandparent()
|
|
grandparent.color = red
|
|
if node == node.Parent.Left && node.Parent == grandparent.Left {
|
|
tree.rotateRight(grandparent)
|
|
} else if node == node.Parent.Right && node.Parent == grandparent.Right {
|
|
tree.rotateLeft(grandparent)
|
|
}
|
|
}
|
|
|
|
func (node *Node) maximumNode() *Node {
|
|
if node == nil {
|
|
return nil
|
|
}
|
|
for node.Right != nil {
|
|
node = node.Right
|
|
}
|
|
return node
|
|
}
|
|
|
|
func (tree *Tree) deleteCase1(node *Node) {
|
|
if node.Parent == nil {
|
|
return
|
|
}
|
|
tree.deleteCase2(node)
|
|
}
|
|
|
|
func (tree *Tree) deleteCase2(node *Node) {
|
|
sibling := node.sibling()
|
|
if nodeColor(sibling) == red {
|
|
node.Parent.color = red
|
|
sibling.color = black
|
|
if node == node.Parent.Left {
|
|
tree.rotateLeft(node.Parent)
|
|
} else {
|
|
tree.rotateRight(node.Parent)
|
|
}
|
|
}
|
|
tree.deleteCase3(node)
|
|
}
|
|
|
|
func (tree *Tree) deleteCase3(node *Node) {
|
|
sibling := node.sibling()
|
|
if nodeColor(node.Parent) == black &&
|
|
nodeColor(sibling) == black &&
|
|
nodeColor(sibling.Left) == black &&
|
|
nodeColor(sibling.Right) == black {
|
|
sibling.color = red
|
|
tree.deleteCase1(node.Parent)
|
|
} else {
|
|
tree.deleteCase4(node)
|
|
}
|
|
}
|
|
|
|
func (tree *Tree) deleteCase4(node *Node) {
|
|
sibling := node.sibling()
|
|
if nodeColor(node.Parent) == red &&
|
|
nodeColor(sibling) == black &&
|
|
nodeColor(sibling.Left) == black &&
|
|
nodeColor(sibling.Right) == black {
|
|
sibling.color = red
|
|
node.Parent.color = black
|
|
} else {
|
|
tree.deleteCase5(node)
|
|
}
|
|
}
|
|
|
|
func (tree *Tree) deleteCase5(node *Node) {
|
|
sibling := node.sibling()
|
|
if node == node.Parent.Left &&
|
|
nodeColor(sibling) == black &&
|
|
nodeColor(sibling.Left) == red &&
|
|
nodeColor(sibling.Right) == black {
|
|
sibling.color = red
|
|
sibling.Left.color = black
|
|
tree.rotateRight(sibling)
|
|
} else if node == node.Parent.Right &&
|
|
nodeColor(sibling) == black &&
|
|
nodeColor(sibling.Right) == red &&
|
|
nodeColor(sibling.Left) == black {
|
|
sibling.color = red
|
|
sibling.Right.color = black
|
|
tree.rotateLeft(sibling)
|
|
}
|
|
tree.deleteCase6(node)
|
|
}
|
|
|
|
func (tree *Tree) deleteCase6(node *Node) {
|
|
sibling := node.sibling()
|
|
sibling.color = nodeColor(node.Parent)
|
|
node.Parent.color = black
|
|
if node == node.Parent.Left && nodeColor(sibling.Right) == red {
|
|
sibling.Right.color = black
|
|
tree.rotateLeft(node.Parent)
|
|
} else if nodeColor(sibling.Left) == red {
|
|
sibling.Left.color = black
|
|
tree.rotateRight(node.Parent)
|
|
}
|
|
}
|
|
|
|
func nodeColor(node *Node) color {
|
|
if node == nil {
|
|
return black
|
|
}
|
|
return node.color
|
|
}
|