Problem Statement (Asked By Google)

A unival tree (short for “universal value tree”) is a tree in which all nodes have the same value.

Given the root of a binary tree, determine how many unival subtrees it contains.

For example, the following tree has 5 unival subtrees:

0
/ \
1 0
/ \
1 0
/ \
1 1

Problem link: https://leetcode.com/problems/count-univalue-subtrees/description/

Disclaimer: Try solving the problem on your own first! Use this solution only as a reference to learn and improve.

---

Solution:

To determine whether a subtree is unival, all its nodes must have the same value. We’ll define a recursive isUnival helper function to check this. A leaf node always qualifies as a unival subtree.

The function to count unival subtrees will use this helper function to evaluate each node.

Approach 1: Recursive O(n^2) Approach:

class TreeNode {
    char value;
    TreeNode left, right;

    TreeNode(char value) {
        this.value = value;
    }
}

public class UnivalSubtrees {

    public static boolean isUnival(TreeNode root, char value) {
        if (root == null) {
            return true;
        }
        return root.value == value && isUnival(root.left, value) && isUnival(root.right, value);
    }

    public static int countUnivalSubtrees(TreeNode root) {
        if (root == null) {
            return 0;
        }
        int left = countUnivalSubtrees(root.left);
        int right = countUnivalSubtrees(root.right);
        return (isUnival(root, root.value) ? 1 : 0) + left + right;
    }

    public static void main(String[] args) {
        TreeNode root = new TreeNode('a');
        root.left = new TreeNode('a');
        root.right = new TreeNode('a');
        root.right.left = new TreeNode('a');
        root.right.right = new TreeNode('a');
        root.right.right.right = new TreeNode('A');

        System.out.println("Unival Subtrees: " + countUnivalSubtrees(root)); // Output: 3
    }
}

Approach 2: Optimized O(n) Approach:

To improve performance, we traverse the tree once, starting from the leaves, and maintain a count of unival subtrees as we return from recursion. This avoids redundant evaluations.

class UnivalSubtrees {

    public static int countUnivalSubtrees(TreeNode root) {
        return helper(root)[0];
    }

    private static int[] helper(TreeNode root) {
        if (root == null) {
            return new int[]{0, 1}; // {count, isUnival (1 for true, 0 for false)}
        }

        int[] left = helper(root.left);
        int[] right = helper(root.right);

        int totalCount = left[0] + right[0];
        boolean isUnival = left[1] == 1 && right[1] == 1;

        if (isUnival) {
            if (root.left != null && root.value != root.left.value) {
                isUnival = false;
            }
            if (root.right != null && root.value != root.right.value) {
                isUnival = false;
            }
        }

        return new int[]{isUnival ? totalCount + 1 : totalCount, isUnival ? 1 : 0};
    }

    public static void main(String[] args) {
        TreeNode root = new TreeNode('a');
        root.left = new TreeNode('c');
        root.right = new TreeNode('b');
        root.right.left = new TreeNode('b');
        root.right.right = new TreeNode('b');
        root.right.right.right = new TreeNode('b');

        System.out.println("Unival Subtrees: " + countUnivalSubtrees(root)); // Output: 5
    }
}

Explanation of Optimized Approach:

1. Base Case: A null node has 0 unival subtrees and is considered a unival subtree itself.
2. Recursive Calculation:
   • Recursively calculate the number of unival subtrees in the left and right subtrees.
   • Check if the current node forms a unival subtree by comparing its value with its children.
   • Add 1 to the total count if the current node is unival.
3. Return:
   • The total count of unival subtrees.
   • A flag indicating whether the current subtree is unival.

Time Complexity:
Both approaches traverse the tree.

• Recursive approach: O(n^2) due to repeated subtree checks.
• Optimized approach: O(n) since each node is visited only once.

Space Complexity: O(h), where h is the height of the tree, for the recursion stack.

---

Did this solution help you understand the concept better? Let me know your thoughts in the comments below! Don’t forget to follow us on Instagram @coderz.py for more DSA tips and solutions to ace your coding interviews.