LeetCode每日一题——1161. 最大层内元素和

题目

给你一个二叉树的根节点 root。设根节点位于二叉树的第 1 层，而根节点的子节点位于第 2 层，依此类推。

请返回层内元素之和 最大 的那几层（可能只有一层）的层号，并返回其中 最小 的那个。

示例

示例 1：

输入：root = [1,7,0,7,-8,null,null]
输出：2
解释：
第 1 层各元素之和为 1，
第 2 层各元素之和为 7 + 0 = 7，
第 3 层各元素之和为 7 + -8 = -1， 所以我们返回第 2 层的层号，它的层内元素之和最大。

示例 2：

输入：root = [989,null,10250,98693,-89388,null,null,null,-32127]
输出：2

提示：

树中的节点数在 [1, 104]范围内
-105 <= Node.val <= 105

思路

1. 使用队列，记录每一层长度、对应层号、该层节点值的总和
2. 每一层计算完比较即可，选出某层节点值总和最大的层号返回

题解

from collections import deque
class Solution:
    def maxLevelSum(self, root: Optional[TreeNode]) -> int:
        if root is None:
            return 
        temp = deque()
        temp.append(root)
        # 参数分别为该层的节点数，最大总和，层数，最大总和的层数
        leng_,max_sum, layer, max_layer = 1, -float('inf'), 0, 0
        while temp:
            sum_, leng_temp = 0, 0
            layer += 1
            for i in range(leng_):
                index = temp.popleft()
                sum_ += index.val
                if index.left:
                    temp.append(index.left)
                    leng_temp += 1
                if index.right:
                    temp.append(index.right)
                    leng_temp += 1
            # 更新下一层的节点数量
            leng_ = leng_temp
            # 更新最大总和即最大总和对应的层数
            if sum_ > max_sum:
                max_sum = sum_
                max_layer = layer
        return max_layer