Tree Diameter

 1 package Tree;
 2 
 3 /*
 4  * The diameter of a tree is the Number of Nodes on the longest path between two leaves in the tree.
 5  */
 6 
 7 public class DiameterBT {
 8     class TreeNode {
 9         int val;
10         TreeNode left, right;
11         public TreeNode(int val) {
12             this.val = val;
13         }
14     }
15     /*
16      * Time : O(n log n)
17      * Space : O(log n)
18      */
19     public int diameter(TreeNode root) {
20         if (root == null) return 0;
21         int dial = diameter(root.left);
22         int diar = diameter(root.right);
23         int hl = height(root.left);
24         int hr = height(root.right);
25         return Math.max(hl + hr + 1, Math.max(dial, diar));
26     }
27     public int height(TreeNode root) {
28         if (root == null) return 0;
29         return Math.max(height(root.left), height(root.right)) + 1;
30     }
31     
32     /*
33      * Time : O(n)
34      */
35     class Height {
36         public int val;
37     }
38     public int diameter(TreeNode root, Height height) {
39         if (root == null) {
40             height.val = 0;
41             return 0;
42         }
43         Height hl = new Height();
44         Height hr = new Height();
45         int dl = diameter(root.left, hl);
46         int dr = diameter(root.right, hr);
47         height.val = Math.max(hl.val, hr.val) + 1;
48         return Math.max(hl.val + hr.val + 1, Math.max(dl, dr));
49     }
50     
51     public static void main(String[] args) {
52         DiameterBT sol = new DiameterBT();
53         TreeNode node0 = sol.new TreeNode(1);
54         TreeNode node1 = sol.new TreeNode(2);
55         TreeNode node2 = sol.new TreeNode(3);
56         TreeNode node3 = sol.new TreeNode(4);
57         TreeNode node4 = sol.new TreeNode(5);
58         node0.left = node1; node0.right = node2;
59         node1.left = node3; node1.right = node4;
60         System.out.println(sol.diameter(node0));
61         System.out.println(sol.diameter(node0, sol.new Height()));
62     }
63 }