Create count-univalue-subtrees.java

Author: maiquynhtruong

count-univalue-subtrees.java

@@ -0,0 +1,50 @@
+/* package whatever; // don't place package name! */
+
+import java.util.*;
+import java.lang.*;
+import java.io.*;
+
+/* Name of the class has to be "Main" only if the class is public. */
+class Ideone
+{
+	static class TreeNode {
+		int val;
+		TreeNode left, right;
+		public TreeNode(int value) {this.val = value;}
+	}
+	public static void main (String[] args) throws java.lang.Exception
+	{
+		TreeNode root = new TreeNode(5);
+		root.left = new TreeNode(1);
+		root.right = new TreeNode(5);
+		root.left.left = new TreeNode(5);
+		root.left.right = new TreeNode(5);
+		root.right.right = new TreeNode(5);
+		System.out.println(countUnivalSubtrees(root));
+	}
+	public static int countUnivalSubtrees(TreeNode root) {
+		if (root == null) return 0;
+		return countUnivals(root);
+	}
+	public static int countUnivals(TreeNode root) {
+		if (root.left == null && root.right == null) {
+			return 1;
+		} else if (root.left != null && root.right != null) {
+			int left = countUnivals(root.left);
+			int right = countUnivals(root.right);
+			if (root.left.val == root.val && root.right.val == root.val) 
+				return 1 + right + left;
+			else return left + right;
+		} else if (root.left != null) {
+			int left = countUnivals(root.left);
+			if (root.left.val == root.val)
+				return left + 1;
+			else return left;
+		} else {
+			int right = countUnivals(root.right);
+			if (root.right.val == root.val)
+				return right + 1;
+			else return right;
+		}
+	}
+}