目标
给你一棵二叉树的根节点 root 和一个正整数 k 。
树中的 层和 是指 同一层 上节点值的总和。
返回树中第 k 大的 层和（不一定不同）。如果树少于 k 层，则返回 -1 。
注意，如果两个节点与根节点的距离相同，则认为它们在同一层。
说明：
- 树中的节点数为 n
- 2 <= n <= 10^5
- 1 <= Node.val <= 10^6
- 1 <= k <= n
思路
这个问题的关键在于遍历的同时记录层数并进行累加。由于是升序排列所以取倒数第k个，即MAX-k。
代码
/**
* @date 2024-02-23 16:06
*/
public class KthLargestLevelSum {
public static class TreeNode {
int val;
TreeNode left;
TreeNode right;
TreeNode() {
}
TreeNode(int val) {
this.val = val;
}
TreeNode(int val, TreeNode left, TreeNode right) {
this.val = val;
this.left = left;
this.right = right;
}
@Override
public String toString() {
return "TreeNode{" +
"val=" + val +
", left=" + left +
", right=" + right +
'}';
}
}
public static void main(String[] args) {
KthLargestLevelSum main = new KthLargestLevelSum();
long res = main.kthLargestLevelSum(new TreeNode(1, new TreeNode(2), new TreeNode(3)), 2);
System.out.println(res);
}
public final int MAX = 100000;
public long[] lacc = new long[MAX];
public int maxlevel = 0;
public long kthLargestLevelSum(TreeNode root, int k) {
traverseSubTree(root, 0);
if (k > maxlevel + 1) {
return -1;
}
Arrays.sort(lacc);
return lacc[MAX-k];
}
public void traverseSubTree(TreeNode subTree, int level) {
maxlevel = Math.max(maxlevel, level);
lacc[level] += subTree.val;
if (subTree.left != null) {
traverseSubTree(subTree.left, level + 1);
}
if (subTree.right != null) {
traverseSubTree(subTree.right, level + 1);
}
}
}Arrays.sort使用的是双轴快排(DualPivotQuicksort)，时间复杂度是O(nlogn)。
性能
最优的算法是快速选择。因为我们并不需要所有的元素都有序，只要保证k位置的左侧都比k小，右侧均比k大即可，而两侧区间内部是不需要排序的。
