Problem Explanation

Using data from a binary tree, we are asked to count the number of uni-value subtrees. A uni-value subtree is a subtree where all nodes have the same value. The goal is, given a binary tree, to return the total number of uni-value subtrees that exist within.

For example, consider the following binary tree:

1
2
3  5
4 / \
51   5
6/ \   \
75  5  5

In this tree, there are 4 uni-value subtrees. They are:

• Subtree rooted at the node with value 1
• Subtree rooted at the first node with value 5
• Subtree rooted at the second node with value 5
• Subtree rooted at the third node with value 5

Approach

To solve this problem, we use the tree traversal technique.

We perform a post-order traversal of the binary tree. This means we first visit the left and right subtree of a node before processing the node itself. So, our helper function isUnival() is recursively called for root->left and root->right. The function returns a boolean to indicate whether the subtree rooted at root is univalued or not.

For each node, we check if the left and right subtrees are both univalued with the same value as the node itself. If they are, we increment our answer and return true. If not, we return false.

This approach has a time complexity of O(n), where n is the total number of nodes in the tree.

Python Solution

1
2python
3class Solution:
4  def countUnivalSubtrees(self, root):
5    self.ans = 0
6    self.is_unival(root, 0)
7    return self.ans
8
9  def is_unival(self, node, val):
10    if not node:
11      return True
12    # Post-order traversal
13    left = self.is_unival(node.left, node.val)
14    right = self.is_unival(node.right, node.val)
15    # If both subtrees are univalued
16    if left and right:
17      self.ans += 1
18      return True if node.val == val else False
19    return False

Java Solution

1
2java
3public class Solution {
4  int count = 0;
5  
6  public int countUnivalSubtrees(TreeNode root) {
7    if(root==null) return 0;
8    isUnival(root,root.val);
9    return count;
10  }
11
12  private boolean isUnival(TreeNode node, int val) {
13    if(node == null) return true;
14    // Check if node.left and node.right are unival trees
15    boolean left = isUnival(node.left, node.val);
16    boolean right = isUnival(node.right, node.val);
17    // Increase count if this node is unival
18    if(left && right) {
19      count++;
20      return node.val == val;
21    }
22    return false;
23  }
24}

C# Solution

1
2csharp
3public class TreeNode {
4    public int val;
5    public TreeNode left;
6    public TreeNode right;
7}
8
9public class Solution {
10    int count;
11
12    public int CountUnivalSubtrees(TreeNode root) {
13        count = 0;
14        helper(root, 0);
15        return count;
16    }
17
18    private bool helper(TreeNode node, int val){
19        if (node == null) return true;
20        bool left = helper(node.left, node.val);
21        bool right = helper(node.right, node.val);
22        if (left && right) {
23            count++;
24            return node.val == val;
25        }
26        return false;
27    }
28}

C++ Solution

1
2cpp
3class Solution {
4public:
5    int countUnivalSubtrees(TreeNode* root) {
6        int ans = 0;
7        isUnival(root, -1, ans);
8        return ans;
9    }
10    
11private:
12    bool isUnival(TreeNode* root, int val, int& ans) {
13        if (!root) return true;
14        // Check if node.left and node.right are unival trees
15        bool left = isUnival(root->left, root->val, ans);
16        bool right = isUnival(root->right, root->val, ans);
17        // Increase count if this node is unival
18        if(left && right) {
19            ans++;
20            return root->val == val;
21        }
22        return false;
23    }
24};

JavaScript Solution

1
2javascript
3class Solution {
4  countUnivalSubtrees(root) {
5    let count = 0;
6    helper(root, 0);
7    return count;
8
9    function isUnival(node, val) {
10      if(node == null) return true;
11      let left = isUnival(node.left, node.val);
12      let right = isUnival(node.right, node.val);
13      if(left && right){
14        count++;
15        return node.val == val;
16      }
17      return false;
18    }
19  }
20}

The JavaScript solution above has a minor bug. The helper function isUnival is not defined inside countUnivalSubtrees and it tries to access count which is local to countUnivalSubtrees. To fix this, we need to define isUnival inside countUnivalSubtrees and also, since this has to be called recursively, it should use the name helper. Here is the corrected JavaScript solution:

1
2javascript
3class Solution {
4  countUnivalSubtrees(root) {
5    let count = 0;
6    
7    function helper(node, val) {
8      if(node == null) return true;
9      let left = helper(node.left, node.val);
10      let right = helper(node.right, node.val);
11      if(left && right){
12        count++;
13        return node.val == val;
14      }
15      return false;
16    }
17    
18    helper(root, 0);
19    return count;
20  }
21}

Now, the JavaScript function countUnivalSubtrees has a helper function helper defined inside it. The helper function checks if the node is Unival (for which it calls itself for the left and right child of the node) similar to the other languages' solutions. The reference to count is valid as helper is defined inside countUnivalSubtrees. The variable count is incremented whenever a Unival subtree is found and the count is returned as a result.

Now, no matter if you are using Python, Java, C#, C++ or JavaScript, you have a solution to count the number of uni-value subtrees in a given binary tree! All solutions rely on the same timeframe of O(n), and provide a good example of the post-order traversal technique. Happy Coding!