Recover Binary Search Tree Python Solution

lcpy gen -n 99   # by problem number
lcpy gen -s recover_binary_search_tree   # by problem name

Input: root = [1,3,null,null,2]
Output: [3,1,null,null,2]
Explanation: 3 cannot be a left child of 1 because 3 > 1. Swapping 1 and 3 makes the BST valid.

Input: root = [3,1,4,null,null,2]
Output: [2,1,4,null,null,3]
Explanation: 2 cannot be in the right subtree of 3 because 2 < 3. Swapping 2 and 3 makes the BST valid.

from leetcode_py import TreeNode


class Solution:
    # Time: O(n)
    # Space: O(1)
    def recover_tree(self, root: TreeNode[int] | None) -> None:
        first: TreeNode[int] | None = None
        second: TreeNode[int] | None = None
        prev: TreeNode[int] | None = None

        current = root
        while current is not None:
            if current.left is None:
                if prev is not None and prev.val > current.val:
                    if first is None:
                        first = prev
                    second = current
                prev = current
                current = current.right
            else:
                predecessor = current.left
                while predecessor.right is not None and predecessor.right is not current:
                    predecessor = predecessor.right
                if predecessor.right is None:
                    predecessor.right = current
                    current = current.left
                else:
                    predecessor.right = None
                    if prev is not None and prev.val > current.val:
                        if first is None:
                            first = prev
                        second = current
                    prev = current
                    current = current.right

        if first is not None and second is not None:
            first.val, second.val = second.val, first.val