one line of code at a time

[leetcode] 1161. Maximum Level Sum of a Binary Tree 파이썬 코드 본문

leetcode

[leetcode] 1161. Maximum Level Sum of a Binary Tree 파이썬 코드

oloc 2024. 8. 16. 09:26

https://leetcode.com/problems/maximum-level-sum-of-a-binary-tree/description/

 

이진 트리가 주어졌을 때 각 레벨별로 합을 구하는데 합이 가장 큰, smallest 레벨을 구하는 것이 문제다. 

레벨 순회가 나오는 순간 bfs (breadth first search)를 떠올렸다.

 

class Solution(object):
    def maxLevelSum(self, root):
        if not root:
            return 0

        queue = []
        queue.append(root)
        result_sum = root.val
        result_level = 1
        curr_level = 0

        while queue:
            temp_sum = 0
            for l in range(len(queue)): # 레벨 순회 
                node = queue.pop(0) 
                temp_sum += node.val
                if node.left:
                    queue.append(node.left)
                if node.right:
                    queue.append(node.right)
            curr_level += 1
            if temp_sum > result_sum:
                result_sum = temp_sum
                result_level = curr_level 

        return result_level

 

그런데 처음에 

node = queue.pop(0) # O
node = queue.pop() # X

 

위에처럼 해야 되는데 아래처럼 해서 코드가 어디가 틀렸는지 찾느라 한참을 헤맸다.

파이썬에서 리스트를 큐처럼 쓰기 위해서는 pop(0)을 해야지, 그냥 pop()을 하면 스택이 되어버린다. 

 

 

deque 쓰기

좀 더 runtime을 줄이고 싶다면, deque을 쓰면 된다. 

class Solution(object):
    def maxLevelSum(self, root):
        if not root:
            return 0

        queue = deque([root])
        result_sum = root.val
        result_level = 1
        curr_level = 0

        while queue:
            temp_sum = 0
            for l in range(len(queue)): 
                node = queue.popleft() 
                temp_sum += node.val
                if node.left:
                    queue.append(node.left)
                if node.right:
                    queue.append(node.right)
            curr_level += 1
            if temp_sum > result_sum:
                result_sum = temp_sum
                result_level = curr_level 

        return result_level