1161. 最大層内元素和

中等樹廣度優先搜索深度優先搜索二叉樹

小貼士：前置題目102. Binary Tree Level Order Traversal（102. 二叉樹的層序遍歷）

思路

使用dfs或是bfs遍歷所有節點，計算每列的數值總和，並找出最大值的列。

如果是bfs，就是層序遍歷。

用queue去紀錄每層節點的內容，並紀錄當前queue的大小size。（代表當前層的節點數量）
從queue當中拿出size次節點，被拿出來的節點都必定在同一層，紀錄到答案當中。
在拿出節點的同時，把節點的左右節點放到queue當中，為下一層的紀錄做準備。

如果是dfs，需要額外紀錄當前所在的位置，因為走訪順序並不按照層。
它會沿著一條路徑一直走到底，直到遇到葉子節點才會回溯，相比bfs來說比較省空間。

※ 需要注意每層的總和可以是負數。

程式碼

1. bfs (queue)

class Solution {
public:
    int maxLevelSum(TreeNode* root) {
        // 題目至少有一個節點，不用判斷 nullptr
        queue<TreeNode*> q;
        q.push(root); // 第一層只有 root 這一個節點
        int best = INT_MIN; // 總和可以是負數
        int res = 1;
        int level = 1;
        while(!q.empty()) {
            int size = q.size();
            int sum = 0;
            while(size--) {
                TreeNode* cur = q.front();
                q.pop();
                sum += cur->val;
                if(cur->left) q.push(cur->left);
                if(cur->right) q.push(cur->right);
            }
            if(best < sum) {
                res = level;
                best = sum;
            }
            ++level;
        }
        return res;
    }
};

使用陣列去實現bfs，在題目有更多變化時，用陣列的方法做能方便根據情況去做修改。比如此題： 103. Binary Tree Zigzag Level Order Traversal 103. 二叉樹的鋸齒形層序遍歷

2. bfs (使用陣列)

class Solution {
private:
    const int MX = 1e4;
public:
    int maxLevelSum(TreeNode* root) {
        vector<TreeNode*> q(MX + 1);
        int left = 0, right = 0;
        q[right++] = root;
        int best = INT_MIN;
        int res = 1;
        int level = 1;
        while(left < right) {
            int size = right - left; // 左閉右開
            int sum = 0;
            while(size--) {
                TreeNode* cur = q[left++];
                sum += cur->val;
                if(cur->left) q[right++] = cur->left;
                if(cur->right) q[right++] = cur->right;
            }
            if(best < sum) {
                res = level;
                best = sum;
            }
            ++level;
        }
        return res;
    }
};

3. dfs (遞迴)

class Solution {
public:
    int maxLevelSum(TreeNode* root) {
        vector<int> sum; // 記住每一層的總和
        auto dfs = [&](this auto&&dfs, TreeNode* node, int level) -> void {
            if(sum.size() <= level) {
                sum.resize(level + 1);
            }
            sum[level] += node->val;
            if(node->left) dfs(node->left, level + 1);
            if(node->right) dfs(node->right, level + 1);
        };
        dfs(root, 0);
        return distance(sum.begin(), ranges::max_element(sum)) + 1; // 用來取代下方區塊, 限 C++20 或以上
        // int res = 1, best = INT_MIN;
        // for(int i = 0; i < sum.size(); ++i) {
        //     if(best < sum[i]) {
        //         best = sum[i];
        //         res = i + 1;
        //     }
        // }
        // return res;
    }
};

4. dfs (迭代)

class Solution {
public:
    int maxLevelSum(TreeNode* root) {
        vector<int> sum; // 記住每一層的總和
        stack<pair<TreeNode*, int>> stk; // {節點, 層數}
        stk.push({root, 0});
        while(!stk.empty()) {
            auto [node, level] = stk.top();
            stk.pop();
            if(sum.size() <= level) {
                sum.resize(level + 1);
            }
            sum[level] += node->val;
            if(node->left) stk.push({node->left, level + 1});
            if(node->right) stk.push({node->right, level + 1});
        }
        return distance(sum.begin(), ranges::max_element(sum)) + 1;
    }
};

複雜度分析

時間複雜度： $O(n)$
空間複雜度： $O(n)$