删除任意元素

Xikl · Nov 8, 2018 · ad17fea · ad17fea
1 parent f3da3ae
commit ad17fea
Showing 1 changed file with 66 additions and 9 deletions.
diff --git a/src/main/java/com/ximo/datastructuresinaction/bst/BinarySearchTree.java b/src/main/java/com/ximo/datastructuresinaction/bst/BinarySearchTree.java
@@ -271,10 +271,7 @@ public E remomveMin() {
     public Node removeMin(Node node) {
         // 左子树为空
         if (node.left == null) {
-            Node rightNode = node.right;
-            node.right = null;
-            size--;
-            return rightNode;
+            return getAndRemoveRight(node);
         }
 
         // 不为空的时候
@@ -302,19 +299,79 @@ public E remomveMax() {
      * @return 删除最小值的 树的根节点
      */
     public Node removeMax(Node node) {
-        // 左子树为空
+        // 右子树为空
         if (node.right == null) {
-            Node leftNode = node.left;
-            node.left = null;
-            size--;
-            return leftNode;
+            return getAndRemoveLeft(node);
         }
 
         // 不为空的时候
         node.right = removeMax(node.right);
         return node;
     }
 
+    public void remove(E element) {
+        root = remove(root, element);
+    }
+
+    private Node remove(Node node, E element) {
+        if (node == null) {
+            return null;
+        }
+        int compareResult = element.compareTo(node.e);
+        if (compareResult < 0) {
+            node.left = remove(node.left, element);
+            return node;
+        } else if (compareResult > 0) {
+            node.right = remove(node.right, element);
+            return node;
+        } else { // 相等的时候
+            // 左子树为空
+            if (node.left == null) {
+                return getAndRemoveRight(node);
+            }
+
+            // 右子树为空
+            if (node.right == null) {
+                return getAndRemoveLeft(node);
+            }
+
+            // 找到后继（替换者）
+            Node successor = minmum(node.right);
+            successor.right = removeMin(node.right);
+            successor.left = node.left;
+
+            // help gc
+            node.left = node.right = null;
+            return successor;
+        }
+    }
+
+    /**
+     * 移除右子树
+     *
+     * @param node 一棵树
+     * @return 返回被删除的右子树
+     */
+    private Node getAndRemoveRight(Node node) {
+        Node rightNode = node.right;
+        node.right = null;
+        size--;
+        return rightNode;
+    }
+
+    /**
+     * 移除左子树
+     *
+     * @param node 一棵树
+     * @return 返回被删除的左子树
+     */
+    private Node getAndRemoveLeft(Node node) {
+        Node leftNode = node.left;
+        node.left = null;
+        size--;
+        return leftNode;
+    }
+
     @Override
     public String toString() {
         StringBuilder res = new StringBuilder();