Merge pull request #4 from geniussportsgroup/develop

Added removeByPos operation

treaps.go

@@ -383,6 +383,36 @@ func (tree *Treap) Remove(key interface{}) interface{} {
 	return retVal.key
 }
 
+func __removePos(rootPtr **Node, i int) *Node {
+
+	root := *rootPtr
+	var retVal *Node
+	if i == root.llink.count {
+		retVal = root
+		*rootPtr = __joinExclusive(&(*rootPtr).llink, &(*rootPtr).rlink)
+		retVal.reset()
+		return retVal
+	} else if i < root.llink.count {
+		retVal = __removePos(&(*rootPtr).llink, i)
+	} else {
+		retVal = __removePos(&(*rootPtr).rlink, i-(root.llink.count+1))
+	}
+
+	root.count--
+
+	return retVal
+}
+
+func (tree *Treap) RemoveByPos(i int) interface{} {
+
+	if i >= tree.Size() {
+		panic(fmt.Sprintf("Invalid position %d", i))
+	}
+
+	retVal := __removePos(tree.rootPtr, i)
+	return retVal.key
+}
+
 // Return the smallest item contained in the tree
 func (tree *Treap) Min() interface{} {
 

treaps_test.go

@@ -584,3 +584,14 @@ func TestTreap_lexicographicCmp(t *testing.T) {
 
 	assert.Equal(t, 1, t5.lexicographicCmp(t1))
 }
+
+func TestTreap_RemoveByPos(t *testing.T) {
+
+	tree := NewTreap(cmpInt, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17)
+
+	assert.Equal(t, 0, tree.RemoveByPos(0))
+	assert.True(t, tree.check())
+
+	assert.Equal(t, 17, tree.RemoveByPos(16))
+	assert.True(t, tree.check())
+}