这题让返回的是第k大的层和,可以使用BFS计算每一层节点的和,然后对这些和从大到小排序,最后取出第k个值即可。
关于二叉树的BFS遍历我们前面也多次讲过,需要使用一个队列来记录每层的节点,当遍历队列中节点的时候需要累加当前层的和,顺便把下一层的节点也添加到队列中。
JAVA:
public long kthLargestLevelSum(TreeNode root, int k) {
List<Long> mList = new ArrayList<>();
Queue<TreeNode> q = new LinkedList<>();// 创建队列
q.offer(root);// 根节点添加到队列中
while (!q.isEmpty()) {
int size = q.size();// 当前层的节点个数
long sum = 0;// 当前层所有节点的和
while (size-- > 0) {
TreeNode cur = q.poll();// 出队
sum += cur.val;// 累加当前层节点的值
// 当前节点的左右子树如果不为空,也添加到队列中
if (cur.left != null) q.offer(cur.left);
if (cur.right != null) q.offer(cur.right);
}
mList.add(sum);// 当前层的和计算完之后添加到list中
}
if (mList.size() < k)// 如果少于k层,返回-1。
return -1;
// 从大到小排序,然后取出第k个元素
Collections.sort(mList, (a, b) -> Long.compare(b, a));
return mList.get(k - 1);
}
C++:
public:
long long kthLargestLevelSum(TreeNode *root, int k) {
vector<long long> mList;
queue<TreeNode *> q;// 创建队列
q.push(root);// 根节点添加到队列中
while (!q.empty()) {
int size = q.size();// 当前层的节点个数
long sum = 0;// 当前层所有节点的和
while (size-- > 0) {
TreeNode *cur = q.front();
q.pop();// 出队
sum += cur->val;// 累加当前层节点的值
// 当前节点的左右子树如果不为空,也添加到队列中
if (cur->left) q.push(cur->left);
if (cur->right) q.push(cur->right);
}
mList.push_back(sum);// 当前层的和计算完之后添加到list中
}
if (mList.size() < k)// 如果少于k层,返回-1。
return -1;
// 从大到小排序,然后取出第k个元素
sort(mList.begin(), mList.end(), greater<long>());
return mList[k - 1];
}
C:
#define n 100000
// 比较函数,用于从大到小排序
int cmp(const void *a, const void *b) {
long long ta = *((long long *) a);
long long tb = *((long long *) b);
if (ta < tb)
return 1;
return -1;
}
long long kthLargestLevelSum(struct TreeNode *root, int k) {
long long *mList = malloc(n * sizeof(long long));
struct TreeNode **q = malloc(n * sizeof(struct TreeNode *));// 创建队列
int left = 0, right = 0;// 左闭右开区间
int count = 0;
q[right++] = root;// 根节点添加到队列中
while (left != right) {
int size = right - left;// 当前层的节点个数
long long sum = 0;// 当前层所有节点的和
while (size-- > 0) {
struct TreeNode *cur = q[left++];// 出队
sum += cur->val;// 累加当前层节点的值
// 当前节点的左右子树如果不为空,也添加到队列中
if (cur->left)
q[right++] = cur->left;
if (cur->right)
q[right++] = cur->right;
}
mList[count++] = sum;// 当前层的和计算完之后添加到mList中
}
if (count < k)// 如果少于k层,返回-1。
return -1;
// 从大到小排序,然后取出第k个元素
qsort(mList, count, sizeof(long long), cmp);// 先排序
long long ans = mList[k - 1];
free(mList);
free(q);
return ans;
}
Python:
def kthLargestLevelSum(self, root: Optional[TreeNode], k: int) -> int:
mList = []
q = deque([root]) # 创建队列
while q:
levelSum = 0 # 当前层所有节点的和
for _ in range(len(q)):
cur = q.popleft() # 出队
levelSum += cur.val # 累加当前层节点的值
# 当前节点的左右子树如果不为空,也添加到队列中
if cur.left: q.append(cur.left)
if cur.right: q.append(cur.right)
mList.append(levelSum) # 当前层的和计算完之后添加到list中
if len(mList) < k: # 如果少于k层,返回-1。
return -1
# 从大到小排序,然后取出第k个元素
mList.sort(reverse=True)
return mList[k - 1]