fixRemoveHeight

avlindex/avlindex.go

@@ -28,7 +28,6 @@ func (n *Node) String() string {
 type Tree struct {
 	root       *Node
 	comparator utils.Comparator
-	count      int
 }
 
 func New(comparator utils.Comparator) *Tree {
@@ -51,6 +50,62 @@ func (avl *Tree) Size() int {
 }
 
 func (avl *Tree) Remove(key interface{}) *Node {
+
+	if n, ok := avl.GetNode(key); ok {
+		if avl.root == n {
+			avl.root = nil
+			return n
+		}
+
+		ls, rs := getChildrenSize(n)
+		if ls == 0 && rs == 0 {
+			p := n.parent
+			p.children[getRelationship(n)] = nil
+			avl.fixRemoveHeight(p)
+			return n
+		}
+
+		var cur *Node
+		if ls > ls {
+			cur = n.children[0]
+			for cur.children[1] != nil {
+				cur = cur.children[1]
+			}
+
+			cleft := cur.children[0]
+			cur.parent.children[getRelationship(cur)] = cleft
+			if cleft != nil {
+				cleft.parent = cur.parent
+			}
+
+		} else {
+			cur = n.children[1]
+			for cur.children[0] != nil {
+				cur = cur.children[0]
+			}
+
+			cright := cur.children[1]
+			cur.parent.children[getRelationship(cur)] = cright
+
+			if cright != nil {
+				cright.parent = cur.parent
+			}
+		}
+
+		cparent := cur.parent
+		// 修改为interface 交换
+		n.value, cur.value = cur.value, n.value
+
+		// 考虑到刚好替换的节点是 被替换节点的孩子节点的时候, 从自身修复高度
+		if cparent == n {
+			avl.fixRemoveHeight(n)
+		} else {
+			avl.fixRemoveHeight(cparent)
+		}
+
+		return cur
+	}
+
 	return nil
 }
 
@@ -434,7 +489,6 @@ func (avl *Tree) fixRemoveHeight(cur *Node) {
 // }
 
 func (avl *Tree) fixPutHeight(cur *Node, lefts, rigths int) {
-	avl.count++
 	if lefts > rigths {
 		l := cur.children[0]
 		llsize, lrsize := getChildrenSize(l)