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.