본문 바로가기
TIL

2583. Kth Largest Sum in a Binary Tree

by 크라00 2024. 10. 22.

- 트리 / 이진 트리 / BFS / 정렬 문제

 

- 문제분석

 

1. 주어진 이진 트리에서 레벨 별 노드 값의 합 중 k번째로 큰 값을 찾는 문제

2. 트리의 각 레벨은 같은 깊이에 있는 모든 노드들로 구성

3. 문제의 목표는 각 레벨의 노드 값의 합을 계산하고, 이 중 k번째로 큰 합을 반환

 

 

- 문제풀이

  1. 레벨 별 합 계산:
    • 이진 트리에서 각 레벨 별 노드 값을 합산하는 가장 쉬운 방법은 **BFS (너비 우선 탐색)**을 사용하는 것입니다.
    • BFS를 사용하면 각 레벨의 노드들을 순차적으로 탐색할 수 있으며, Queue를 이용하여 쉽게 구현할 수 있습니다.
    • Queue에 루트 노드를 추가하고, 각 레벨마다 Queue에서 노드를 하나씩 꺼내어 자식들을 다시 큐에 추가하면서 현재 레벨의 노드 합을 계산합니다.
  2. k번째 큰 합 계산:
    • 최대값을 기준으로 정렬된 값 중에서 k번째를 찾는 것이므로, 효율적인 방법으로 PriorityQueue (Python에서는 heapq)를 사용하여 힙 자료구조를 이용해 레벨 합을 관리합니다.
    • **최소 힙(Min-Heap)**을 유지하며 최대 k개의 레벨 합만 저장하여 힙의 크기를 k로 유지합니다. 이렇게 하면 시간과 메모리를 효율적으로 사용할 수 있습니다.
  3. 시간 복잡도:
    • BFS 탐색에 의해 모든 노드를 한 번씩 방문하므로, 트리의 총 노드 개수를 n이라 할 때 **O(n)**의 시간 복잡도가 필요합니다.
    • 레벨 합의 수가 많아야 O(log k)의 시간 복잡도를 가지므로, 최종적인 시간 복잡도는 **O(n + n log k)**가 됩니다.

접근 방법

  1. BFS로 레벨 별 합 계산:
    • Queue에 루트 노드를 추가하고, 큐가 빌 때까지 반복하면서 각 레벨을 순차적으로 탐색합니다.
    • 각 레벨의 노드 값을 모두 더한 뒤, 그 값을 리스트에 추가합니다.
  2. 최소 힙을 사용하여 k번째 큰 합 계산:
    • 각 레벨 합을 리스트에 저장한 후, k개의 가장 큰 값을 찾기 위해 최소 힙을 사용합니다.
    • 힙의 크기를 k로 유지하면서 가장 작은 값들을 제거하고, 최종적으로 힙의 루트에 남아 있는 값이 k번째 큰 값이 됩니다.

 

 

> java

 

class Solution {
    public long kthLargestLevelSum(TreeNode root, int k) {
        if (root == null) {
            return -1;
        }

        Queue<TreeNode> queue = new LinkedList<>();
        queue.offer(root);
        List<Long> levelSums = new ArrayList<>();

        while(!queue.isEmpty()) {

            int size = queue.size();
            long levelSum = 0;
            for (int i = 0; i < size; i++) {
                TreeNode node = queue.poll();
                levelSum += node.val;

                if (node.left != null) {
                    queue.offer(node.left);
                }
                if (node.right != null) {
                    queue.offer(node.right);
                }
            }
            levelSums.add(levelSum);
        }

        PriorityQueue<Long> minHeap = new PriorityQueue<>(k);
        for (long sum : levelSums) {
            minHeap.offer(sum);
            if (minHeap.size() > k) {
                minHeap.poll();
            }
        }

        return minHeap.size() == k ? minHeap.peek() : -1;
    }
}

 

 

> python

 

from typing import Optional
from collections import deque
import heapq

class Solution:
    def kthLargestLevelSum(self, root: Optional[TreeNode], k: int) -> int:
        if root is None :
            return -1
        
        queue = deque([root])
        level_sums = []

        while queue :
            size = len(queue)
            level_sum = 0
            for _ in range(size) :
                node = queue.popleft()
                level_sum += node.val

                if node.left is not None :
                    queue.append(node.left)
                if node.right is not None :
                    queue.append(node.right)
            level_sums.append(level_sum)


        # 최소 힙을 사용하여 k번째 큰 값을 찾기
        if len(level_sums) < k:
            return -1

        min_heap = []
        for sum_value in level_sums:
            heapq.heappush(min_heap, sum_value)
            if len(min_heap) > k:
                heapq.heappop(min_heap)

        return heapq.heappop(min_heap)