11
11
12
12
from __future__ import annotations
13
13
14
+ from collections .abc import Iterator
15
+
14
16
15
17
class Node :
16
18
"""
@@ -33,24 +35,24 @@ class BinaryTreeNodeSum:
33
35
12 8 0
34
36
35
37
>>> tree = Node(10)
36
- >>> BinaryTreeNodeSum(tree).node_sum( )
38
+ >>> sum( BinaryTreeNodeSum(tree))
37
39
10
38
40
39
41
>>> tree.left = Node(5)
40
- >>> BinaryTreeNodeSum(tree).node_sum( )
42
+ >>> sum( BinaryTreeNodeSum(tree))
41
43
15
42
44
43
45
>>> tree.right = Node(-3)
44
- >>> BinaryTreeNodeSum(tree).node_sum( )
46
+ >>> sum( BinaryTreeNodeSum(tree))
45
47
12
46
48
47
49
>>> tree.left.left = Node(12)
48
- >>> BinaryTreeNodeSum(tree).node_sum( )
50
+ >>> sum( BinaryTreeNodeSum(tree))
49
51
24
50
52
51
53
>>> tree.right.left = Node(8)
52
54
>>> tree.right.right = Node(0)
53
- >>> BinaryTreeNodeSum(tree).node_sum( )
55
+ >>> sum( BinaryTreeNodeSum(tree))
54
56
32
55
57
"""
56
58
@@ -64,8 +66,8 @@ def depth_first_search(self, node: Node | None) -> int:
64
66
self .depth_first_search (node .left ) + self .depth_first_search (node .right )
65
67
)
66
68
67
- def node_sum (self ) -> int :
68
- return self .depth_first_search (self .tree )
69
+ def __iter__ (self ) -> Iterator [ int ] :
70
+ yield self .depth_first_search (self .tree )
69
71
70
72
71
73
if __name__ == "__main__" :
0 commit comments