leetcode: Binary Tree Level Order Traversal (python / 파이썬)

요즘에는 개발자들이 회사에 취업하기 위해서 사용한다는 리트코드 (leetcode) 사이트의 문제들을 풀어보고 있습니다. 이제 막 시작한 초보자라 힘든 점이 많아서 이렇게 글로 남겨서 복습하려고 합니다.

 

문제: Binary Tree Level Order Traversal
등급: MEDIUM

글 적는 내 모습

<코드>

어떻게 하면 각 level 별로 리스트를 출력할 수 있을까... 한 30분간 고민했던 것 같다. 그런데 생각을 못했다. ㅠㅠ

 

우선 기존의 preorder, postorder 등은 잘 할 수 있을 거라고 생각된다. 그런데, level 별 리스트를 어떻게 하지?

 

count을 사용해서 각 level의 count를 하나씩 늘려가면서 해볼까 생각해봤는데 그 level를 언제 +1 하게 만들 것인가라는 질문에 대답을 하지 못하였다.

 

너무나도 아이디어가 떠오르지 않은 나머지, discussion을 살짝 훔쳐봤다. 거기서 set를 이용하는 것이 보였다. 한 번 시도해봐야겠다.

 

아, 모르겠다.

 

결국에는 discussion을 다 보게 되었다. 역시, 강호는 넓고 고수는 많다고 하였나. 왜 이렇게 잘 짜여진 코드들이 많은지...

 

그런데 읽어도 이해 안되는 것들이 꽤 되었다. 우선, 처음 이해해서 적은 코드는 이와 같다:

if not root:
            return []
        
        ans, each_level = [], [root]
        while each_level:
            ans.append([node.val for node in each_level])
            temp = []
            for node in each_level:
                temp.extend([node.left, node.right]) 
            level = [leaf_node for leaf_node in temp if leaf_node]
            
        return ans

그런데 memory limit exceeded가 뜬다고....? 왜...? 어째서...? 똑같은 코드로 이름을 바꾸면 되던데, 그래서 포기했다.

 

하지만, 논리를 설명하자면,

기본적으로 root 없으면 empty list를 반환한다.

 

그 후, ans 리스트와 each_level 리스트 값 할당을 한다. each_level 리스트는 각 레벨의 node를 모두 포함하게 될 것이다. 처음에는 일단 당연히 [root] 밖에 없다.

 

그렇게 each_level을 while loop을 돌려본다 (자연스럽게 마지막 level에 node가 없다면 멈출 것). 우선 ans에 each_level 리스트에 있는 node들의 값을 리스트로 넣으며, temp []를 만들고 for loop을 통해 각 node.left와 node.right을 extend()해준다. 그런 다음 temp에 저장된 node들을 null이 아니라면 level 리스트에 넣어준다. 

 

이를 반복하다보면 답이 구해진다. 

# Definition for a binary tree node.
# class TreeNode(object):
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution(object):
    def levelOrder(self, root):
        ans, eachlevel = [], [root]
        
        while root and eachlevel:
            ans.append([node.val for node in eachlevel])
            lrpair = [(node.left, node.right) for node in eachlevel]
            eachlevel = [leaf for landr in lrpair for leaf in landr if leaf]
        return ans

그리고 같은 논리의 짧은 코드

 

<오류, 힘들었던 점>

- 계속 generator object한테는 node.left와 node.right가 없다고 해서 뭐지? 했다. 알고 보니, 

each_level  = [leaf_node for leaf_node...] 가 아닌 each_level.append(leaf_node for leaf_node ..)를 했다. 이러니 당연히 generator object가 append 되지.... 하아...

 

<총평 외 comment>

 

Happy coding!

Binary Tree Level Order Traversal - LeetCode