Delete Node in a BST Python Solution

lcpy gen -n 450   # by problem number
lcpy gen -s delete_node_in_a_bst   # by problem name

Input: root = [5,3,6,2,4,null,7], key = 3
Output: [5,4,6,2,null,null,7]
Explanation: One valid answer is [5,4,6,2,null,null,7]; [5,2,6,null,4,null,7] is also accepted.

Input: root = [5,3,6,2,4,null,7], key = 0
Output: [5,3,6,2,4,null,7]
Explanation: The tree does not contain a node with value = 0.

Input: root = [], key = 0
Output: []

from leetcode_py import TreeNode


class Solution:
    # Time: O(h) where h is the height of the tree
    # Space: O(h) recursion stack
    def delete_node(self, root: TreeNode[int] | None, key: int) -> TreeNode[int] | None:
        if root is None:
            return None

        if key < root.val:
            root.left = self.delete_node(root.left, key)
        elif key > root.val:
            root.right = self.delete_node(root.right, key)
        else:
            # Node found: handle deletion by child count.
            if root.left is None:
                return root.right
            if root.right is None:
                return root.left
            # Two children: replace value with inorder successor, delete successor.
            successor = root.right
            while successor.left is not None:
                successor = successor.left
            root.val = successor.val
            root.right = self.delete_node(root.right, successor.val)
        return root