more on bst

1 files changed, 49 insertions, 40 deletions

diff --git a/set/tree.go b/set/tree.go
index 8233b0e..819f83c 100644
--- a/set/tree.go
+++ b/set/tree.go
@@ -19,63 +19,72 @@ func (t *Tree) Empty() bool {
 	return t.root == nil
 }
 
-func (t *Tree) set(n *node, key, val int) {
-	switch {
-	case key < n.key:
-		if n.left == nil {
-			n.left = &node{key, val, nil, nil}
-			return
-		}
-		t.set(n.left, key, val)
-
-	case key > n.key:
-		if n.right == nil {
-			n.right = &node{key, val, nil, nil}
-			return
-		}
-		t.set(n.right, key, val)
-	default:
-		// Val is already in the tree
-	}
-}
-
 func (t *Tree) Set(key, val int) {
 	if t.Empty() {
 		t.root = &node{key, val, nil, nil}
 		return
 	}
-	t.set(t.root, key, val)
+	ptr, _, err := t.search(&t.root, key)
+	if err != nil {
+		// key already in the tree as node found
+		return
+	}
+
+	*ptr = &node{key, val, nil, nil}
 }
 
-func (t *Tree) get(n *node, key int) (int, error) {
-	if n == nil {
-		return 0, NotFound
+func (t *Tree) Get(key int) (int, error) {
+	_, n, err := t.search(&t.root, key)
+	return n.val, err
+}
+
+func (t *Tree) Del(key int) (int, error) {
+	ptr, n, err := t.search(&t.root, key)
+	if err != nil {
+		return n.val, err
 	}
 
+	// Case 1: n is leaf node
+	// Case 2: n has one child
+	// Case 3: n has two childs
+
 	switch {
-	case key < n.key:
-		return t.get(n.left, key)
-	case key > n.key:
-		return t.get(n.right, key)
-	default:
+	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 childs!
+		return 0, NotImplemented
 	}
 }
 
-func (t *Tree) Get(key int) (int, error) {
-	return t.get(t.root, key)
+func (t *Tree) min(ptr **node) (**node, *node, error) {
+	return nil, nil, NotImplemented
 }
 
-func (t *Tree) Del(key int) (int, error) {
-	if t.root == nil {
-		return 0, NotFound
+func (t *Tree) search(ptr **node, key int) (**node, *node, error) {
+	n := *ptr
+	if n == nil {
+		return ptr, nil, NotFound
 	}
 
-	if t.root.key == key {
-		val := t.root.val
-
-		return val, nil
+	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
 	}
-
-	return 0, NotImplemented
 }