LeetCode 250 同值子樹計數：遞歸後序遍歷解法

題目描述

給定一棵二元樹，統計所有「同值子樹」的數量。同值子樹定義：該子樹所有節點值均相同。

範例：

    5
   / \
  1   5
 / \   \
5   5   5

回傳 4（三個葉子節點 5、右子樹 5-5、以及根節點 5 的右子樹）。

演算法分析

採用後序遍歷（Post-order），自底向上判斷每棵子樹是否為同值子樹：

• 空節點視為同值（回傳 0，不計數）
• 葉子節點必為同值子樹（計數 +1，回傳 1）
• 內部節點：若左右子樹皆為同值（回傳 1）且值等於當前節點值，則當前子樹也是同值（計數 +1，回傳 1）
• 若單側子樹為同值且值相等，另一側為空（回傳 0），亦為同值（計數 +1，回傳 1）
• 其他情況回傳 -1 表示非同值

C++ 實作

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    int countUnivalSubtrees(TreeNode* root) {
        int res = 0;
        dfs(root, res);
        return res;
    }

private:
    // 回傳：1=同值子樹, 0=空節點, -1=非同值
    int dfs(TreeNode* root, int& res) {
        if (!root) return 0;
        if (!root->left && !root->right) {
            ++res;
            return 1;
        }

        int l = dfs(root->left, res);
        int r = dfs(root->right, res);

        // 兩側皆為同值子樹
        if (l == 1 && r == 1) {
            if (root->val == root->left->val && root->val == root->right->val) {
                ++res;
                return 1;
            }
        }
        // 任一側非同值
        if (l == -1 || r == -1) return -1;

        // 單側同值且值相等
        if (l != 0 && root->val == root->left->val) {
            ++res;
            return 1;
        }
        if (r != 0 && root->val == root->right->val) {
            ++res;
            return 1;
        }
        return -1;
    }
};

複雜度

• 時間：O(N)，每個節點訪問一次
• 空間：O(H)，遞歸堆疊深度，H 為樹高（最壞 O(N)）

圖示說明