Java实现LeetCode2583.二叉树中的第K大层和

题目描述

给你一棵二叉树的根节点 root 和一个正整数 k 。

树中的 层和 是指 同一层 上节点值的总和。

返回树中第 k 大的层和（不一定不同）。如果树少于 k 层，则返回 -1 。

注意 ，如果两个节点与根节点的距离相同，则认为它们在同一层。

思路

先开始看漏了一个字，看成了第K层的和了。
第K大层和，可以维护一个优先队列。
如果最终队列的size小于k，说明不足k层返回-1，否则返回队列头部。

代码

class Solution {
    public long kthLargestLevelSum(TreeNode root, int k) {
        Deque<TreeNode> dq = new LinkedList();
        dq.add( root );
        PriorityQueue<Long> pq = new PriorityQueue();
        while( !dq.isEmpty() ){
            int size = dq.size();
            long cnt = 0;
            while( size > 0 ){
                TreeNode tmp = dq.poll();
                cnt += tmp.val; 
                if( tmp.left!=null ) dq.add( tmp.left );
                if( tmp.right!=null ) dq.add( tmp.right );
                size--;
            }
            if( pq.size()<k ) pq.add( cnt );
            else{
                if( cnt > pq.peek() ){
                    pq.poll();
                    pq.add( cnt );
                }
            }
        }
        if( pq.size()<k ) return -1;
        return pq.peek();
    }
}