Delete Node in a BST

1package pkg5leetcode.official75;2 3/*4 * Delete Node in a BST | LC 4505 * APPROACH: BST delete with successor for two-child case.6 * COMPLEXITY: Time O(h), Space O(h)7 */8public class official75_LC450DeleteNodeInABST {9    /** Same shape as pkg5leetcode/common/TreeNode.java (nested for single-file runs). */10 11    static class TreeNode {12        int val;13        TreeNode left, right;14        TreeNode(int val) { this.val = val; }15    }16 17    static TreeNode deleteNode(TreeNode root, int key) {18        if (root == null) return null;19        if (key < root.val) root.left = deleteNode(root.left, key);20        else if (key > root.val) root.right = deleteNode(root.right, key);21        else {22            if (root.left == null) return root.right;23            if (root.right == null) return root.left;24            TreeNode succ = root.right;25            while (succ.left != null) succ = succ.left;26            root.val = succ.val;27            root.right = deleteNode(root.right, succ.val);28        }29        return root;30    }31 32    static boolean contains(TreeNode root, int key) {33        while (root != null) {34            if (root.val == key) return true;35            root = key < root.val ? root.left : root.right;36        }37        return false;38    }39 40    public static void main(String[] args) {41        TreeNode root = new TreeNode(5);42        root.left = new TreeNode(3); root.right = new TreeNode(6);43        root.left.left = new TreeNode(2); root.left.right = new TreeNode(4);44        root.right.right = new TreeNode(7);45        deleteNode(root, 3);46        check(!contains(root, 3) && contains(root, 2), "case1");47        TreeNode r2 = new TreeNode(5); r2.left = new TreeNode(3); r2.right = new TreeNode(6);48        r2.left.left = new TreeNode(2); r2.left.right = new TreeNode(4);49        deleteNode(r2, 5);50        check(r2.val == 6 && contains(r2, 4), "case2");51        System.out.println("all tests passed");52    }53 54    static void check(boolean cond, String name) {55        if (!cond) throw new AssertionError("FAILED: " + name);56        System.out.println("  PASS " + name);57    }58}