diff options
Diffstat (limited to 'search/bst.go')
| -rw-r--r-- | search/bst.go | 171 |
1 files changed, 171 insertions, 0 deletions
diff --git a/search/bst.go b/search/bst.go new file mode 100644 index 0000000..0eecb9e --- /dev/null +++ b/search/bst.go @@ -0,0 +1,171 @@ +package set + +import "fmt" + +type node struct { + key int + val int + left *node + right *node +} + +func (n *node) String() string { + recurse := func(n *node) string { + if n == nil { + return "" + } + return n.String() + } + + return fmt.Sprintf("node{%d:%d,%s,%s}", + n.key, + n.val, + recurse(n.left), + recurse(n.right)) +} + +type BST struct { + root *node +} + +func NewBST() *BST { + return &BST{} +} + +func (t *BST) String() string { + if t.Empty() { + return "BST{}" + } + + return fmt.Sprintf("BST{%s}", t.root) +} + +func (t *BST) Empty() bool { + return t.root == nil +} + +func (t *BST) Set(key, val int) { + if t.Empty() { + t.root = &node{key, val, nil, nil} + return + } + ptr, _, err := t.search(&t.root, key) + switch err { + case nil: + // key already in the tree + return + case NotFound: + *ptr = &node{key, val, nil, nil} + return + default: + panic(err) + } +} + +func (t *BST) Get(key int) (int, error) { + _, n, err := t.search(&t.root, key) + if err != nil { + return 0, err + } + return n.val, nil +} + +func (t *BST) Del(key int) (int, error) { + ptr, n, err := t.search(&t.root, key) + if err != nil { + return 0, err + } + + // Case 1: n is leaf node + // Case 2: n has one child + // Case 3: n has two childs + + switch { + case n.left == nil: + if n.right == nil { + // I am a leaf node + *ptr = nil + return n.val, nil + } + // I have a right child + *ptr = n.right + return n.val, nil + + case n.right == nil: + // I have a left child + *ptr = n.left + return n.val, nil + default: + // I have two children! + + o, err := t.deleteMin(&n.right) + if err != nil { + return 0, err + } + + o.left = n.left + o.right = n.right + *ptr = o + + return n.val, nil + } +} + +func (t *BST) search(ptr **node, key int) (**node, *node, error) { + n := *ptr + if n == nil { + return ptr, nil, NotFound + } + + switch { + case key < n.key: + return t.search(&n.left, key) + case n.key < key: + return t.search(&n.right, key) + default: + return ptr, n, nil + } +} + +func (t *BST) deleteMin(ptr **node) (*node, error) { + ptr, n, err := t.min(ptr) + if err != nil { + return nil, err + } + + *ptr = n.right + n.right = nil + return n, nil +} + +func (t *BST) min(ptr **node) (**node, *node, error) { + n := *ptr + if n == nil { + return nil, nil, NotFound + } + + for { + if n.left == nil { + return ptr, n, nil + } + ptr = &n.left + n = *ptr + } +} + +/* +func (t *BST) max(ptr **node) (**node, *node, error) { + n := *ptr + if n == nil { + return nil, nil, NotFound + } + + for { + if n.right == nil { + return ptr, n, nil + } + ptr = &n.right + n = *ptr + } +} +*/ |
