Merge pull request #7 from preetapan/master

Implement DeletePrefix for feature parity with go-immutable-radix, pl…

Author: armon

radix.go

@@ -292,6 +292,53 @@ DELETE:
 	return leaf.val, true
 }
 
+// DeletePrefix is used to delete the subtree under a prefix
+// Returns how many nodes were deleted
+// Use this to delete large subtrees efficiently
+func (t *Tree) DeletePrefix(s string) int {
+	return t.deletePrefix(nil, t.root, s)
+}
+
+// delete does a recursive deletion
+func (t *Tree) deletePrefix(parent, n *node, prefix string) int {
+	// Check for key exhaustion
+	if len(prefix) == 0 {
+		// Remove the leaf node
+		subTreeSize := 0
+		//recursively walk from all edges of the node to be deleted
+		recursiveWalk(n, func(s string, v interface{}) bool {
+			subTreeSize++
+			return false
+		})
+		if n.isLeaf() {
+			n.leaf = nil
+		}
+		n.edges = nil // deletes the entire subtree
+
+		// Check if we should merge the parent's other child
+		if parent != nil && parent != t.root && len(parent.edges) == 1 && !parent.isLeaf() {
+			parent.mergeChild()
+		}
+		t.size -= subTreeSize
+		return subTreeSize
+	}
+
+	// Look for an edge
+	label := prefix[0]
+	child := n.getEdge(label)
+	if child == nil || (!strings.HasPrefix(child.prefix, prefix) && !strings.HasPrefix(prefix, child.prefix)) {
+		return 0
+	}
+
+	// Consume the search prefix
+	if len(child.prefix) > len(prefix) {
+		prefix = prefix[len(prefix):]
+	} else {
+		prefix = prefix[len(child.prefix):]
+	}
+	return t.deletePrefix(n, child, prefix)
+}
+
 func (n *node) mergeChild() {
 	e := n.edges[0]
 	child := e.node

radix_test.go

@@ -104,6 +104,46 @@ func TestDelete(t *testing.T) {
 	}
 }
 
+func TestDeletePrefix(t *testing.T) {
+	type exp struct {
+		inp[] string
+		prefix string
+		out[] string
+		numDeleted int
+	}
+
+	cases := []exp{
+		{[]string{"", "A", "AB", "ABC", "R", "S"}, "A", []string{"", "R", "S"},  3},
+		{[]string{"", "A", "AB", "ABC", "R", "S"}, "ABC", []string{"", "A", "AB", "R", "S"},  1},
+		{[]string{"", "A", "AB", "ABC", "R", "S"}, "", []string{},  6},
+		{[]string{"", "A", "AB", "ABC", "R", "S"}, "S", []string{"", "A", "AB", "ABC", "R"},  1},
+		{[]string{"", "A", "AB", "ABC", "R", "S"}, "SS", []string{"", "A", "AB", "ABC", "R", "S"},  0},
+	}
+
+	for _, test := range cases {
+		r := New()
+		for _, ss := range test.inp {
+			r.Insert(ss, true)
+		}
+
+		deleted := r.DeletePrefix(test.prefix)
+		if deleted != test.numDeleted {
+			t.Fatalf("Bad delete, expected %v to be deleted but got %v", test.numDeleted, deleted)
+		}
+
+		out := []string{}
+		fn := func(s string, v interface{}) bool {
+			out = append(out, s)
+			return false
+		}
+		r.Walk(fn)
+
+		if !reflect.DeepEqual(out, test.out) {
+			t.Fatalf("mis-match: %v %v", out, test.out)
+		}
+	}
+}
+
 func TestLongestPrefix(t *testing.T) {
 	r := New()