2583. Kth Largest Sum in a Binary Tree

Problem : https://leetcode.com/problems/kth-largest-sum-in-a-binary-tree

 

Difficulty : Medium

 

Status : Solved

 

Time : 00:04:22

 

---

 

문제 설명

 

이진 트리의 루트 root와 양의 정수 k가 주어진다.

 

트리의 레벨합은 같은 레벨에 있는 모든 노드의 값의 합을 의미한다.

 

k번째로 큰 레벨합을 구하여라. 트리 레벨이 k보다 적다면 -1을 반환하라.

 

두 노드는 루트 노드와 같은 거리에 있다면 같은 레벨이다.

---

 

풀이

 

여러 방법이 있겠지만, 나는 BFS를 먼저 시도했다.

 

트리에서 BFS를 수행하면, 한 사이클을 돌 때마다 한 레벨의 노드를 전부 방문할 수 있다. 즉 O(N) 시간복잡도 내에 추가적인 알고리즘 없이 레벨순으로 모든 노드를 방문할 수 있다. 따라서 레벨합을 매우 쉽게 구할 수 있고, 이 레벨합 중 K번째 레벨합을 구하면 된다(이는 보통 O(NlogN) 시간복잡도가 걸릴 것이다)

 

풀이 코드

# 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 kthLargestLevelSum(self, root: Optional[TreeNode], k: int) -> int:
        nums = []
        curr = [root]

        while curr :
            res = 0
            tmp = []
            for c in curr :
                res += c.val
                if c.left is not None :
                    tmp.append(c.left)
                if c.right is not None :
                    tmp.append(c.right)
            heappush(nums, -res)
            curr = tmp

        if len(nums) < k :
            return -1
        
        for _ in range(k-1) :
            heappop(nums)
        return -heappop(nums)