tpop/datastructures.go

127 lines
2.2 KiB
Go

package main
import (
"fmt"
"io"
)
type TreeNode interface {
Left() TreeNode
Right() TreeNode
Print(io.Writer)
}
type BTreeNode struct {
left TreeNode
right TreeNode
}
type StringyBTreeNode struct {
k string
v interface{}
BTreeNode
}
func NewStringyBTree(k string, v interface{}) *StringyBTreeNode {
return &StringyBTreeNode{
k: k,
v: v,
}
}
func (sbt *StringyBTreeNode) Left() TreeNode {
return sbt.left
}
func (sbt *StringyBTreeNode) Right() TreeNode {
return sbt.right
}
func (sbt *StringyBTreeNode) Print(w io.Writer) {
fmt.Fprintf(w, "k: %s v: %v, ", sbt.k, sbt.v)
l, lok := sbt.left.(*StringyBTreeNode)
r, rok := sbt.right.(*StringyBTreeNode)
if !rok {
fmt.Fprint(w, "r: _, ")
} else {
fmt.Fprintf(w, "r: %s, ", r.k)
}
if !lok {
fmt.Fprint(w, "l: _")
} else {
fmt.Fprintf(w, "l: %s", l.k)
}
fmt.Fprint(w, "\n")
if sbt.left != nil {
sbt.left.Print(w)
}
if sbt.right != nil {
sbt.right.Print(w)
}
}
func (sbt *StringyBTreeNode) Insert(k string, v interface{}) *StringyBTreeNode {
if sbt.k == k {
sbt.v = v // set new value
return sbt
}
newNode := &StringyBTreeNode{k: k, v: v}
if sbt.k > k {
if sbt.left == nil {
sbt.left = newNode
} else {
sbt.left.(*StringyBTreeNode).Insert(k, v)
}
}
if sbt.k < k {
if sbt.right == nil {
sbt.right = newNode
} else {
sbt.right.(*StringyBTreeNode).Insert(k, v)
}
}
return sbt
}
func (sbt *StringyBTreeNode) Get(k string) interface{} {
r := sbt.getNode(k)
if r == nil {
return nil
}
return r.v
}
func (sbt *StringyBTreeNode) getNode(k string) *StringyBTreeNode {
if sbt.k == k {
return sbt
}
if sbt.left != nil {
r := sbt.left.(*StringyBTreeNode).getNode(k)
if r != nil {
return r
}
}
if sbt.right != nil {
return sbt.right.(*StringyBTreeNode).getNode(k)
}
return nil // not found
}
// swap the rightmost node to the node we're removing
func (sbt *StringyBTreeNode) Delete(k string) *StringyBTreeNode {
parent := sbt
rightmost := sbt
for {
if rightmost.right == nil {
break
} else {
parent, rightmost = rightmost, parent.right.(*StringyBTreeNode)
}
}
delete := sbt.getNode(k)
delete.k = rightmost.k
delete.v = rightmost.v
parent.right = nil
return sbt
}