LeetCode: 99. Recover Binary Search Tree

1. 초안

코드가 뭔가 복잡하다.

코드도 줄일겸 메모리를 O(1)으로 풀어보자.

# 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:
    def recoverTree(self, root: Optional[TreeNode]) -> None:
        """
        Do not return anything, modify root in-place instead.
        """
        
        self.arr = []
        self.helper(root)
        
        left, right = None, None
        for i in range(len(self.arr)-1):
            if self.arr[i][0] > self.arr[i+1][0]:
                left = i
                break
       
        i = left + 1
        min_val = pow(2, 31) + 1
        while i <= len(self.arr) - 1 and self.arr[i][0] < self.arr[left][0]:
            if self.arr[i][0] < min_val:
                min_val = self.arr[i][0]
                right = i
            i += 1
        
        left = self.arr[left][1]
        right = self.arr[right][1]
        temp = left.val
        left.val = right.val
        right.val = temp
        
    def helper(self, cur):
        if cur.left != None:
            self.helper(cur.left)
        self.arr.append([cur.val, cur])
        if cur.right != None:
            self.helper(cur.right)

 

2. 메모리를 O(1)으로 줄인 방법

# 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:
    def recoverTree(self, root: Optional[TreeNode]) -> None:
        """
        Do not return anything, modify root in-place instead.
        """
        INF = pow(2, 31) + 1
        # self.reg1: fisrt error value
        # self.reg2: second error value
        self.reg1, self.reg2 = -INF, INF
        # found_flag: True if an error has been found
        self.found_flag = False
        # swap node1's & node2's value
        self.node1, self.node2 = None, None
        
        self.dfs(root)
        
        temp = self.node1.val
        self.node1.val = self.node2.val
        self.node2.val = temp
        
        
    def dfs(self, cur):
        if cur.left != None:
            self.dfs(cur.left)
            
        val = cur.val
        # if no error until now
        if not self.found_flag:
            # valid
            if val > self.reg1:
                self.reg1 = val
                self.node1 = cur
            # error
            else:
                self.found_flag = True
        if self.found_flag:
            if val < self.reg1:
                if val < self.reg2:
                    self.reg2 = val
                    self.node2 = cur
            else:
                return
        
        if cur.right != None:
            self.dfs(cur.right)