add deletion, make Get work after rebalance

This commit is contained in:
David 2021-07-25 20:37:11 -04:00
parent 508460da67
commit f9d933e33c
2 changed files with 78 additions and 10 deletions

View File

@ -38,7 +38,20 @@ func (sbt *StringyBTreeNode) Right() TreeNode {
} }
func (sbt *StringyBTreeNode) Print(w io.Writer) { func (sbt *StringyBTreeNode) Print(w io.Writer) {
fmt.Printf("%+v\n", sbt) l, lok := sbt.left.(*StringyBTreeNode)
r, rok := sbt.right.(*StringyBTreeNode)
fmt.Fprintf(w, "k: %s v: %v, ", sbt.k, sbt.v)
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 { if sbt.left != nil {
sbt.left.Print(w) sbt.left.Print(w)
} }
@ -71,18 +84,39 @@ func (sbt *StringyBTreeNode) Insert(k string, v interface{}) *StringyBTreeNode {
} }
func (sbt *StringyBTreeNode) Get(k string) interface{} { 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 { if sbt.k == k {
return sbt.v return sbt
} }
if sbt.k > k { if sbt.left != nil {
if sbt.left != nil { return sbt.left.(*StringyBTreeNode).getNode(k)
return sbt.left.(*StringyBTreeNode).Get(k)
}
} }
if sbt.k < k { if sbt.right != nil {
if sbt.right != nil { return sbt.right.(*StringyBTreeNode).getNode(k)
return sbt.right.(*StringyBTreeNode).Get(k)
}
} }
return nil // not found return nil // not found
} }
func (sbt *StringyBTreeNode) rightmost(parent *StringyBTreeNode) (*StringyBTreeNode, *StringyBTreeNode) {
if sbt.right == nil {
return parent, sbt
}
return sbt.right.(*StringyBTreeNode).rightmost(sbt)
}
// swap the rightmost node to the node we're removing
func (sbt *StringyBTreeNode) Delete(k string) *StringyBTreeNode {
parent, rightmost := sbt.rightmost(nil)
delete := sbt.getNode(k)
delete.k = rightmost.k
delete.v = rightmost.v
parent.right = nil
return sbt
}

View File

@ -1,6 +1,7 @@
package main package main
import ( import (
"os"
"testing" "testing"
) )
@ -27,6 +28,39 @@ func TestBuildingTree(t *testing.T) {
} }
} }
func TestDeletingFromTree(t *testing.T) {
tree := NewStringyBTree("foo", 1).Insert("bar", 2).Insert("baz", 3).Insert("quuz", 4)
tree.Delete("bar")
if tree.Get("foo") != 1 {
t.Fail()
}
if tree.Get("bar") != nil {
t.Fail()
}
if tree.Get("baz") != 3 {
t.Fail()
}
if tree.Get("quuz") != 4 {
t.Fail()
}
tree.Print(os.Stdout)
tree = NewStringyBTree("foo", 1).Insert("bar", 2).Insert("baz", 3).Insert("quuz", 4)
tree.Delete("foo")
if tree.Get("foo") != nil {
t.Fail()
}
if tree.Get("bar") != 2 {
t.Fail()
}
if tree.Get("baz") != 3 {
t.Fail()
}
if tree.Get("quuz") != 4 {
t.Fail()
}
}
var r TreeNode var r TreeNode
func BenchmarkInserts(b *testing.B) { func BenchmarkInserts(b *testing.B) {