Arafatk/DataViz

// 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"
    "strconv"

    "github.com/Arafatk/Dataviz/trees"
    "github.com/Arafatk/Dataviz/utils"
)

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
    nodeIndex int
    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}
}

// 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{}) {
    var insertedNode *Node
    if tree.Root == nil {
        tree.Root = &Node{Key: key, Value: value, color: red}
        insertedNode = tree.Root
    } 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 = &Node{Key: key, Value: value, color: red}
                    insertedNode = node.Left
                    loop = false
                } else {
                    node = node.Left
                }
            case compare > 0:
                if node.Right == nil {
                    node.Right = &Node{Key: key, Value: value, color: red}
                    insertedNode = node.Right
                    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
}

// 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)
    }
}

// Visualizer makes a visual image demonstrating the Red Black Tree data structure
// using dot language and Graphviz. It first producs a dot string corresponding
// to the Red Black Tree and then runs graphviz to output the resulting image to a file.
func (tree *Tree) Visualizer(fileName string) bool {
    dotString := "digraph graphname{" // Initializing dot string
    var colorArray []color            // Array containing colors of all nodes
    it := tree.Iterator()
    NodeIndex := 0 // An index used to mark different nodes so that node connections are easily represented in the graph
    for i := 0; it.Next(); i++ {
        it.node.nodeIndex = NodeIndex
        NodeIndex++
        colorArray = append(colorArray, it.NodeColor())
    }
    NilNodes := NodeIndex // The nil leaf nodes of the tree
    it = tree.Iterator()
    for i := 0; it.Next(); i++ { // Making all node connections
        if it.node.Left != nil { // Left Child
            dotString += (strconv.Itoa(it.node.nodeIndex) + " -> " + strconv.Itoa(it.node.Left.nodeIndex) + ";")
        } else {
            dotString += (strconv.Itoa(it.node.nodeIndex) + " -> " + strconv.Itoa(NilNodes) + ";")
            NilNodes++
        }
        if it.node.Right != nil { // Right Child
            dotString += (strconv.Itoa(it.node.nodeIndex) + " -> " + strconv.Itoa(it.node.Right.nodeIndex) + ";")
        } else {
            dotString += (strconv.Itoa(it.node.nodeIndex) + " -> " + strconv.Itoa(NilNodes) + ";")
            NilNodes++
        }
    }
    stringValues := []string{} // Putting all the elements of a stack to a string array.
    values := tree.Values()
    Keys := tree.Keys()

    for i := 0; i < NodeIndex; i++ { // Labeling individual nodes with colors and adding data labels
        stringValues = append(stringValues, fmt.Sprintf("%v", Keys[i]))
        stringValues = append(stringValues, fmt.Sprintf("%v", values[i]))
        if colorArray[i] {
            dotString += (strconv.Itoa(i) + "[color=black, style=filled, fillcolor = black, fontcolor=white,label=\"" + stringValues[len(stringValues)-2] + "->" + stringValues[len(stringValues)-1] + "\"];")
        } else {
            dotString += (strconv.Itoa(i) + "[color=red, style=filled, fillcolor = red, fontcolor=white,label=\"" + stringValues[len(stringValues)-2] + "->" + stringValues[len(stringValues)-1] + "\"];")

        }
    }

    for i := NodeIndex; i < NilNodes; i++ { // Adding labels for nil nodes
        dotString += (strconv.Itoa(i) + "[color=coral, style=\"rounded,filled\", shape=box, fillcolor = coral, fontcolor=white,label=Nil];")
    }
    dotString += "}" // Finishing the DotString
    return utils.WriteDotStringToPng(fileName, dotString)
}

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
}