Counting subtrees where all nodes have the same value

Question

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:

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


class DailyCodingProblem8 {

    public static void main(String args[]) {

        BinaryTree tree = new BinaryTree();
        tree.root = new Node(0);
        tree.root.left = new Node(1);
        tree.root.right = new Node(0);
        tree.root.right.left = new Node(1);
        tree.root.right.right = new Node(0);
        tree.root.right.left.left = new Node(1);
        tree.root.right.left.right = new Node(1);
        int res = tree.countUnivalTrees(tree.root);
        System.out.println(res);

        /*
              5
             / \
            4   5
           / \   \
          4   4   5

        */

        tree = new BinaryTree();
        tree.root = new Node(5);
        tree.root.left = new Node(4);
        tree.root.right = new Node(5);
        tree.root.left.left= new Node(4);
        tree.root.left.right= new Node(4);
        tree.root.right.right = new Node(5);
        res = tree.countUnivalTrees(tree.root);
        System.out.println(res);

    }

}

class Node {
    public int value;
    public Node left, right;

    Node(int value) {
        this.value = value;
        this.left = this.right = null;
    }
}

class BinaryTree {
    Node root;

    int countUnivalTrees(Node root) {
        if (root == null) {
            return 0;
        }
        int count = countUnivalTrees(root.left) + countUnivalTrees(root.right);
        if (root.left != null && root.value != root.left.value) {
            return count;
        }

        if (root.right != null && root.value != root.right.value) {
            return count;
        }

        // if whole tree is unival tree
        return count + 1;
    }

}

What is the best way to supply binary tree as input? Should I be creating a insert method and insert nodes? Will the interviewer feel that I am deviating from the actual problem if I do so?

Answer 1

Allow Node take left and right, and BinaryTree take a node.

new BinaryTree(
    new Node(
        0,
        new Node(1),
        new Node(
            0,
            new Node(
                1,
                new Node(1)
                new Node(1)
            ),
            new Node(0),
        )
    )
)

Takes a fair amount of lines, but clearly shows the shape of the tree and removes the need for the messy tree.root.right.left.right usage.