파이썬 알고리즘 인터뷰 - 이진 트리 반전

문제

이진 트리가 주어지면 root트리를 반전 하고 루트 를 반환합니다.

입력

입력과 결과

root = [4,2,7,1,3,6,9]

출력

[4,7,2,9,6,3,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

def list_to_tree(lst, idx):
    parent = None
    if idx < len(lst):
        value = lst[idx]
        if value is not None:
            parent = TreeNode(value)
            parent.left = list_to_tree(lst, idx*2+1)
            parent.right = list_to_tree(lst, idx*2+2)
    return parent

def tree_len(root):
    cnt = 0
    stack = [root]
    while stack:
        node = stack.pop()
        if node is not None:
            cnt += 1
            stack.append(node.left)
            stack.append(node.right)
    return cnt

def tree_to_list(root, lst, idx=0):
    node = root
    if node:
        lst[idx] = node.val
    if node.left:
        tree_to_list(node.left, lst, idx * 2 + 1)
    if node.right:
        tree_to_list(node.right, lst, idx * 2 + 2)

    return lst

def invertTree(root: [TreeNode]) -> [TreeNode]:
    stack = [root]
    while stack:
        node = stack.pop()
        if node:
            node.left, node.right = node.right, node.left

            stack.append(node.left)
            stack.append(node.right)
    return root

r = list_to_tree([4,2,7,1,3,6,9], 0)
lst = tree_to_list(r,[0]*tree_len(r))
print(lst)
invertR = invertTree(r)
print(tree_to_list(invertR,[0]*tree_len(invertR)))

이 문제를 내 코드에서 확인 해보려면 배열을 트리로 바꿔주고, 그 결과가 TreeNode로 반환되는데, 그걸 또 배열로 바꿔주어야 확인 할 수 있어서 직접 list_to_tree, tree_to_list를 구현해서 확인 해봤다.
직접 구현해서 사용해 보니까 훨씬 이해가 쉬워지고 자신감이 생겼다. ^0^

https://leetcode.com/problems/invert-binary-tree/

Invert Binary Tree - LeetCode