# Time: O(N)
# Space: O(N)

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    items = swaps = None
    
    def recoverTree(self, root: Optional[TreeNode]) -> None:
        """
        Do not return anything, modify root in-place instead.
        """
        # Get the almost-sorted tree as an array
        self.items = self.inorder(root)
        
        # Find the elements that need to be swapped
        self.swaps = self.identify_swaps(self.items)
        
        # Walk the tree and fix the matching values
        self.walk_fix(root, 0)
        
    
    def inorder(self, node):
        """
        Return list of elements in almost sorted order (BST property
        implies that in-order traversal produces sorted list)
        """
        
        return self.inorder(node.left) + [node.val] + self.inorder(node.right) if node is not None else []
        
        
    def identify_swaps(self, items):
        # We need to swap two values so keep two variables
        x = y = None

        for i in range(len(items) - 1):
            # Problem states that we at most need to swap one item
            # at index i to another item at index j to get the sorted
            # list. If this is the case, the larger value to swap will
            # be seen first and later on, the smaller value will be seen
            
            # e.g. 1 4 3 2
            # and  2 1
            if items[i] < items[i + 1]:
                pass
            else:
                # Condition fails when we see (4, 3) in first e.g and (2, 1)
                # in second example. We immediately assign i+1th value to cover
                # a case like the second example where swaps are right next to each
                # in last iteration
                y = items[i + 1]

                # In (2, 1)'s case, immediately goes to this block and iteration stops
                # when iteration condition becomes False.
                #
                # In first example's case, we get to (4, 3) in which case:
                #
                # y = items[i + 1] = 3
                # x = items[i] = 4
                #
                # We keep looping until we hit (3, 2) and this time, it's the lower value
                # so y is reassigned to 2. x is already set previously to 4 so now we can
                # break and return the values to swap.
                if x is None:
                    x = items[i]
                else:
                    break

        return {
            x: y,
            y: x
        }
        
        
    def walk_fix(self, node, i):
        """
        Does in-order traversal of BST and if node.val matches
        a key in self.swaps, replaces it with the swapped value
        """
        if node is None: return

        self.walk_fix(node.left, i)

        if node.val in self.swaps:
            node.val = self.swaps[node.val]

        i += 1

        self.walk_fix(node.right, i)