1
+ '''
2
+ Diameter of Binary Tree
3
+
4
+ Given a binary tree, you need to compute the length of the diameter of the tree.
5
+ The diameter of a binary tree is the length of the longest path between any two nodes in a tree.
6
+ This path may or may not pass through the root.
7
+ Note: The length of path between two nodes is represented by the number of nodes.
8
+
9
+ Input: 3
10
+ / \
11
+ 1 4
12
+ \ \
13
+ 2 5
14
+ /
15
+ 7
16
+ Output: 6
17
+ Output Explanation: [7, 2, 1, 3, 4, 5] is the diameter of the binary tree.
18
+
19
+ Input: 5
20
+ / \
21
+ 3 6
22
+ / \
23
+ 2 4
24
+ / \
25
+ 1 8
26
+ Output: 5
27
+ Output Explanation: [1, 2, 3, 4, 8] is the diameter of the binary tree.
28
+
29
+ =========================================
30
+ Traverse the tree and keep/return information about the longest/max branch and longest/max diameter.
31
+ Time Complexity: O(N)
32
+ Space Complexity: O(N) , because of the recursion stack (but this is if the tree is one branch), O(LogN) if the tree is balanced.
33
+ '''
34
+
35
+
36
+ ############
37
+ # Solution #
38
+ ############
39
+
40
+ # import TreeNode class from tree_helpers.py
41
+ from tree_helpers import TreeNode
42
+
43
+ def diameter (root ):
44
+ return find_diameter (root )[1 ]
45
+
46
+ def find_diameter (root ):
47
+ ''' returns (max branch length, max diameter) '''
48
+ if not root :
49
+ return 0 , 0
50
+
51
+ # traverse left and right subtrees
52
+ left , right = find_diameter (root .left ), find_diameter (root .right )
53
+
54
+ # return the max branch from the left and right subtrees plus the current node
55
+ # and find the max diameter till now (using the current node and the max left and right subtree branches)
56
+ return max (left [0 ], right [0 ]) + 1 , max (left [1 ], right [1 ], left [0 ] + right [0 ] + 1 )
57
+
58
+
59
+ ###########
60
+ # Testing #
61
+ ###########
62
+
63
+ # Test 1
64
+ # Correct result => 6
65
+ tree = TreeNode (3 , TreeNode (1 , None , TreeNode (2 , TreeNode (7 ))), TreeNode (4 , None , TreeNode (5 )))
66
+ print (diameter (tree ))
67
+
68
+ # Test 2
69
+ # Correct result => 5
70
+ tree = TreeNode (5 , TreeNode (3 , TreeNode (2 , TreeNode (1 )), TreeNode (4 , None , TreeNode (8 ))), TreeNode (6 ))
71
+ print (diameter (tree ))
0 commit comments