Python/data_structures/binary_tree/diameter_of_binary_tree.py at master · twilight-debugger/Python · GitHub

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
72
73
"""
The diameter/width of a tree is defined as the number of nodes on the longest path
between two end nodes.
"""

from __future__ importannotations

fromdataclassesimportdataclass


@dataclass
classNode:
data: int
left: Node|None=None
right: Node|None=None

defdepth(self) ->int:
"""
 >>> root = Node(1)
 >>> root.depth()
 1
 >>> root.left = Node(2)
 >>> root.depth()
 2
 >>> root.left.depth()
 1
 >>> root.right = Node(3)
 >>> root.depth()
 2
 """
left_depth=self.left.depth() ifself.leftelse0
right_depth=self.right.depth() ifself.rightelse0
returnmax(left_depth, right_depth) +1

defdiameter(self) ->int:
"""
 >>> root = Node(1)
 >>> root.diameter()
 1
 >>> root.left = Node(2)
 >>> root.diameter()
 2
 >>> root.left.diameter()
 1
 >>> root.right = Node(3)
 >>> root.diameter()
 3
 """
left_depth=self.left.depth() ifself.leftelse0
right_depth=self.right.depth() ifself.rightelse0
returnleft_depth+right_depth+1


if__name__=="__main__":
fromdoctestimporttestmod

testmod()
root=Node(1)
root.left=Node(2)
root.right=Node(3)
root.left.left=Node(4)
root.left.right=Node(5)
r"""
 Constructed binary tree is
 1
 / \
 2 3
 / \
 4 5
 """
print(f"{root.diameter() =}") # 4
print(f"{root.left.diameter() =}") # 3
print(f"{root.right.diameter() =}") # 1