2583. 二叉树中的第 K 大层和 【BFS + 排序】

53 阅读1分钟

2583. 二叉树中的第 K 大层和

层序遍历依次计算记录层和 + 排序 O(nlogn)O(n)\lgroup O(n\log n) 、O(n)\rgroup

# 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:
        # 计算 每一层 的层和
        # 返回 第 k 大的

        q = deque([root])
        lis = []
        while q:
            n = len(q)
            s = 0
            for _ in range(n):            
                node = q.popleft()
                s += node.val 
                if node.left:
                    q.append(node.left)
                if node.right:
                    q.append(node.right)

            lis.append(s)

        if len(lis) < k:
            return -1 

        lis.sort()
        return lis[-k]

image.png

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
public:
    long long kthLargestLevelSum(TreeNode* root, int k) {
        vector<long long> arr;
        vector<TreeNode*> q = {root};
        while (!q.empty()){
            long long s = 0;
            vector<TreeNode*> nxt;
            for (auto node : q){// 用 中间临时数组  无需弹出
                s += node->val;
                if (node->left) nxt.emplace_back(node->left);
                if (node->right) nxt.emplace_back(node->right);
            }
            arr.emplace_back(s);
            q = move(nxt); //  赋值
        }
        if (k > arr.size())  return -1;
        ranges::sort(arr);
        return arr[arr.size() - k];
    }
};