Counting number of unival subtrees in a binary tree in Python

By Vedant Keshav Jadhav

This is a time efficient Python program to count the number of unival subtrees in a binary tree. This algorithm takes O(n) time, where n is number of nodes in the tree.

A universal value tree, also called as unival tree, is a tree data structure where all the nodes in the tree are having the same value as the root. An example of it is:

 

The problem statement is to count the number of unival subtrees in a given binary tree. In the binary tree, every node has at most 2 children. An example is:

In the above tree, the number of unival subtrees is 5.

To count the number of unival subtrees, we have two methods, brute force, which is slow, and an efficient approach.

 

Approach 1:(Brute Force)  In this approach, for every node, we check if the right and left subtrees of that node have same value as root. We know that if they are unival, all the values in the left subtree must be equal to the value in root, and same should for the right subtree. If satisfied we increment the count, else we check for nect nodes. The psuedo code for this method is as follows:

// Function takes a root node and returns the count of unival subtrees in the tree 
countUnivals(root):
    if(!root)
        return 0

    if(isUnival(root))
        return 1   
 
    return countUnivals(root.left) + countUnivals(root.right) + 1


//Checking if a subtree is unival or not
isUnival_subtree(root,key): 
    if(root==NULL)
        return true 

    return (root.data==key and isUnival_subtree(root.left,key) && isUnival_subtree(root.right,key)) 


//For a particular node, check if it is unival
isUnival(root): 
    if(root==NULL) 
        return true
    key=root.data
   
    return isUnival_subtree(root,key)

Time complexity - Here, for every node, we check whether the values in it's right and left subtrees are equal to the node data. So for k^th node, we will check (k - 1) nodes. So the time complexity for n nodes will be -

T(n) - T(n - 1) + T(n - 2) + T(n - 3) + ... = O(n²)

 

Approach 2:  This is a more time efficient algorithm. In this approach, we use bottom-up approach, by traversing the complete binary tree once, and using recursion. We use the knowledge that every leaf node is a unival subtree of it's parent node, and that a given tree is unival if it's right subtree and left subtree are unival and their value is equal to the ndoe value. Psuedo code:

//function to count number of unival subtrees
count_Univals(root):
    static variable count = 0

    //base cases
    if(!root)
        return True
    if(root.left==NULL and root.right==NULL)
        count = count + 1
        return True

    //recursion
    left_val = count_Univals(root.left)
    right_val = count_Univals(root.right)
    if(left_vals and right_vals)
        if(!root.left and root.data != root.right.data)
            return False
        if(!root.right and root.data != root.left.data)
            return False
        if(root.rigth.data != root.data or root.left.data != root.data)
            return False
        
        count = count + 1
        return True

Time complexity - In this approach, we traverse the entire tree only once in a recursive manner. So the time complexity is:

T(n) = O(n)

 

Python program - This algorithm is implemented in the Python program. The input to the code is in string format, with comma-separated values, following the given template:

Input - [ Numbers with comma separating two numbers ] Here, null values are represented by "null" word.

For example, the input for the above trees having 0s and 1s is -

[ 0, 1, null, null, 0, 1, 1, null, null, 1, null, null, 0]

 

Output - Output is an integer value, denoting the number of unival subtrees in the given binary tree.

For the above input, the output will be 5.