Universal tree - CoolMind

Problem
Solution

Problem

This problem was asked by Google.

A unival tree (which stands for “universal value”) is a tree where all nodes under it have the same value.

Given the root to a binary tree, count the number of unival subtrees.

For example, the following tree has 5 unival subtrees:

Solution

Assume an empty tree has no subtrees
All leaf nodes are unival subtrees
If a node is not a unival subtrees, then all its parents, parents of parents are not.

struct TreeNode {
    int va;
    TreeNode* left;
    TreeNode* right;
};

int count_unival_substrees(TreeNode* root) {

    int count = 0;
    function<bool(TreeNode*)> traverse = [&](TreeNode* node) {
        if (!node) return true;
        if (!node->left && !node->right) {
            // early stop here
            ++count;
            return true;
        }
        bool left = traverse(node->left);
        bool right = traverse(node->right);
        if (!left || !right) return false;
        if (node->left && node->left-val != node->val) {
            return false;
        }
        if (node->right && node->right-val != node->val) {
            return false;
        }
        ++count;
        return true;
    };
    traverse(node);
    return count;
}