Count Unival Subtrees in a Binary Tree

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.

Recent Posts

Longest Absolute Path in File System Representation

Find the length of the longest absolute path to a file within the abstracted file…

3 days ago

Efficient Order Log Storage

You manage an e-commerce website and need to keep track of the last N order…

2 weeks ago

Select a Random Element from a Stream

You are given a stream of elements that is too large to fit into memory.…

3 weeks ago

Estimate π Using Monte Carlo Method

The formula for the area of a circle is given by πr². Use the Monte…

1 month ago

Longest Substring with K Distinct Characters

Given an integer k and a string s, write a function to determine the length…

1 month ago

Staircase Climbing Ways

There is a staircase with N steps, and you can ascend either 1 step or…

1 month ago