Link: https://leetcode.cn/problems/count-univalue-subtrees/

Question

difficulty: mid
adj diff: 3

Given the root of a binary tree, return the number of uni-value subtrees.

A uni-value subtree means all nodes of the subtree have the same value.



Example 1:

Input: root = [5,1,5,5,5,null,5]
Output: 4

Example 2:

Input: root = []
Output: 0

Example 3:

Input: root = [5,5,5,5,5,null,5]
Output: 6



Constraints:

    The number of the node in the tree will be in the range [0, 1000].
    -1000 <= Node.val <= 1000

标准 tree 的 dfs()。

我的代码不是很巧妙，因为 return value 搞出来了 2种特殊情况。

下面的更新版代码更好。

Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
class Solution {

    Map<Integer, Integer> map = new HashMap<Integer, Integer>();
    int total = 0;

    public int countUnivalSubtrees(TreeNode root) {
        if (root != null) {
            dfs(root);
        }
        return total;
    }

    private int dfs(TreeNode node) {
        if (node == null) return -10000;
        int leftVal = dfs(node.left);
        int rightVal = dfs(node.right);

        if (leftVal + rightVal == -20000) {
            total++;
            map.put(node.val, map.getOrDefault(node.val, 0) + 1);
            return node.val;
        } else if (rightVal == -10000 && node.val == leftVal) {
            total++;
            map.put(node.val, map.getOrDefault(node.val, 0) + 1);
            return node.val;
        } else if (leftVal == -10000 && node.val == rightVal) {
            total++;
            map.put(node.val, map.getOrDefault(node.val, 0) + 1);
            return node.val;
        } else if (node.val == leftVal && node.val == rightVal) {
            total++;
            map.put(node.val, map.getOrDefault(node.val, 0) + 1);
            return node.val;
        }
        return 10000;
    }
}

可以 simply 成如下短代码。

（相当于把 return value 当作 parameter 传到下一级 dfs function）

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
class Solution {
    
    int total = 0;

    public int countUnivalSubtrees(TreeNode root) {
        dfs(root, 0);
        return total;
    }

    private boolean dfs(TreeNode node, int val) {
        if (node == null) return true;

        boolean leftCheck = dfs(node.left, node.val);
        boolean rightCheck = dfs(node.right, node.val);
        if (leftCheck && rightCheck) {
            total++;
        } else {
            return false;
        }
        // [IMPORTANT] Can't simply say: 
        // if (!dfs(node.left, node.val) || !dfs(node.right, node.val)) return false;
        // beceause if (left consider) is checked, (right consider) won't check

        return node.val == val;
    }
}