2583. Kth Largest Sum in a Binary Tree

题目 #

You are given the root of a binary tree and a positive integer k.

The level sum in the tree is the sum of the values of the nodes that are on the same level.

Return the kth largest level sum in the tree (not necessarily distinct). If there are fewer than k levels in the tree, return -1.

Note that two nodes are on the same level if they have the same distance from the root.

Example 1:

Input: root = [5,8,9,2,1,3,7,4,6], k = 2
Output: 13
Explanation: The level sums are the following:
- Level 1: 5.
- Level 2: 8 + 9 = 17.
- Level 3: 2 + 1 + 3 + 7 = 13.
- Level 4: 4 + 6 = 10.
The 2nd largest level sum is 13.

Example 2:

Input: root = [1,2,null,3], k = 1
Output: 3
Explanation: The largest level sum is 3.

Constraints:

  • The number of nodes in the tree is n.
  • $2 <= n <= 10^5$
  • $1 <= Node.val <= 10^6$
  • 1 <= k <= n

思路1 bfs+小根堆 #

分析 #

  • bfs遍历每一层
  • 最小堆找第k大

代码 # 

 1type LittleHeap []int64
 2
 3func (h *LittleHeap) Len() int { return len(*h) }
 4
 5// less必须满足当Less(i, j)和Less(j, i)都为false，则两个索引对应的元素相等
 6// 为true，i向栈顶移动；为false，j向栈顶移动
 7func (h *LittleHeap) Less(i, j int) bool { return (*h)[i] < (*h)[j] }
 8func (h *LittleHeap) Swap(i, j int)      { (*h)[i], (*h)[j] = (*h)[j], (*h)[i] }
 9func (h *LittleHeap) Push(x interface{}) {
10	*h = append(*h, x.(int64))
11}
12
13func (h *LittleHeap) Pop() interface{} {
14	x := (*h)[len(*h)-1]
15	*h = (*h)[:len(*h)-1]
16	return x
17}
18
19/**
20 * Definition for a binary tree node.
21 * type TreeNode struct {
22 *     Val int
23 *     Left *TreeNode
24 *     Right *TreeNode
25 * }
26 */
27func kthLargestLevelSum(root *TreeNode, k int) int64 {
28	q := make([]*TreeNode, 0, 1)
29	tmp := make([]*TreeNode, 0, 1)
30	bHeap := make(LittleHeap, 0, k)
31
32	q = append(q, root)
33	for len(q) > 0 {
34		tmp, q = q, tmp
35		var s int64 = 0
36		q = q[:0]
37		for _, node := range tmp {
38			s += int64(node.Val)
39			if node.Left != nil {
40				q = append(q, node.Left)
41			}
42			if node.Right != nil {
43				q = append(q, node.Right)
44			}
45		}
46		if len(bHeap) < k {
47			heap.Push(&bHeap, s)
48		} else if s > bHeap[0] {
49			heap.Pop(&bHeap)
50			heap.Push(&bHeap, s)
51		}
52	}
53	if len(bHeap) < k {
54		return -1
55	}
56	return bHeap[0]
57}