forked from MTrajK/coding-problems
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathdiameter_of_binary_tree.py
More file actions
71 lines (56 loc) · 1.94 KB
/
diameter_of_binary_tree.py
File metadata and controls
71 lines (56 loc) · 1.94 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
'''
Diameter of Binary Tree
Given a binary tree, you need to compute the length of the diameter of the tree.
The diameter of a binary tree is the length of the longest path between any two nodes in a tree.
This path may or may not pass through the root.
Note: The length of path between two nodes is represented by the number of nodes.
Input: 3
/ \
1 4
\ \
2 5
/
7
Output: 6
Output Explanation: [7, 2, 1, 3, 4, 5] is the diameter of the binary tree.
Input: 5
/ \
3 6
/ \
2 4
/ \
1 8
Output: 5
Output Explanation: [1, 2, 3, 4, 8] is the diameter of the binary tree.
=========================================
Traverse the tree and keep/return information about the longest/max branch and longest/max diameter.
Time Complexity: O(N)
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.
'''
############
# Solution #
############
# import TreeNode class from tree_helpers.py
from tree_helpers import TreeNode
def diameter(root):
return find_diameter(root)[1]
def find_diameter(root):
''' returns (max branch length, max diameter) '''
if not root:
return 0, 0
# traverse left and right subtrees
left, right = find_diameter(root.left), find_diameter(root.right)
# return the max branch from the left and right subtrees plus the current node
# and find the max diameter till now (using the current node and the max left and right subtree branches)
return max(left[0], right[0]) + 1, max(left[1], right[1], left[0] + right[0] + 1)
###########
# Testing #
###########
# Test 1
# Correct result => 6
tree = TreeNode(3, TreeNode(1, None, TreeNode(2, TreeNode(7))), TreeNode(4, None, TreeNode(5)))
print(diameter(tree))
# Test 2
# Correct result => 5
tree = TreeNode(5, TreeNode(3, TreeNode(2, TreeNode(1)), TreeNode(4, None, TreeNode(8))), TreeNode(6))
print(diameter(tree))