LeetCode/Top Interview 150

114. Flatten Binary Tree to Linked List

hyunkookim 2024. 12. 13. 17:38

114. Flatten Binary Tree to Linked List

 

# 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 flatten(self, root: Optional[TreeNode]) -> None:
        """
        Do not return anything, modify root in-place instead.
        """
        # preorder: 왼쪽쭉 내려가서. 올라오면서, right
        if not root: # root 가 없으면
            return []

        stack = [root] # 초기 스텍에 루트 노드 추가
        
        while stack:
            cur = stack.pop() # 스텍에서 현재 노드 갖고오기

            # 오른쪽, 왼쪽 자식 추가
            # 왼쪽이 먼저 나올려면, 오른쪽이 먼저 들어가야 함
            if cur.right:
                stack.append(cur.right)
            if cur.left:
                stack.append(cur.left)

            cur.left = None # 현재 노드의 왼쪽을 제거하고
            # 오른쪽을 다음노드로 설정
            if stack:
                cur.right = stack[-1]

 

https://youtu.be/rKnD7rLT0lI?si=I5hfPBRDbTDoXT8q

 

Follow up: Can you flatten the tree in-place (with O(1) extra space)?

# 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 flatten(self, root: Optional[TreeNode]) -> None:
        """
        이진 트리를 연결 리스트 형태로 평탄화 (in-place 변환).
        """

        # DFS를 사용해 트리를 평탄화하는 헬퍼 함수.
        # 현재 노드(root)의 마지막 노드(tail)를 반환.
        def dfs(root):
            if not root:  # 종료 조건: 노드가 None이면 아무 작업도 하지 않고 반환.
                return None

            # 현재 서브트리의 루트 상태:
            #
            #       root
            #      /    \
            #  root.left  root.right

            # 왼쪽과 오른쪽 서브트리를 재귀적으로 평탄화.
            leftTail = dfs(root.left)  # 왼쪽 서브트리를 평탄화하고 마지막 노드(tail)를 가져옴.
            rightTail = dfs(root.right)  # 오른쪽 서브트리를 평탄화하고 마지막 노드(tail)를 가져옴.

            # 왼쪽 서브트리가 존재할 경우 처리.
            if root.left:
                # 왼쪽 서브트리의 마지막 노드(leftTail)의 오른쪽에 오른쪽 서브트리를 연결.
                leftTail.right = root.right
                
                # 현재 노드의 오른쪽에 왼쪽 서브트리를 배치.
                root.right = root.left

                # 왼쪽 서브트리를 제거.
                root.left = None

                # 현재 트리 구조 (왼쪽 서브트리를 오른쪽으로 이동):
                #       root
                #        \
                #         root.left (이전 왼쪽 서브트리)
                #           \
                #            root.right (이전 오른쪽 서브트리)

            # 현재 서브트리의 마지막 노드를 반환.
            # 오른쪽 서브트리가 있으면 그 마지막 노드를 반환.
            # 그렇지 않으면 왼쪽 서브트리의 마지막 노드를 반환.
            # 둘 다 없으면 현재 노드를 반환.
            last = rightTail or leftTail or root
            return last

        # 트리가 다음과 같은 형태로 주어졌다고 가정:
        #
        #     1
        #    / \
        #   2   5
        #  / \    \
        # 3   4    6
        #

        dfs(root)  # 재귀 호출로 평탄화 수행

        # 최종 결과 트리:
        #
        # 1
        #  \
        #   2
        #    \
        #     3
        #      \
        #       4
        #        \
        #         5
        #          \
        #           6