Skip to content

Commit c9078c7

Browse files
committed
refactor: Change replace method with __iter__ overriding (TheAlgorithms#7020)
1 parent 413f4f9 commit c9078c7

File tree

1 file changed

+9
-7
lines changed

1 file changed

+9
-7
lines changed

data_structures/binary_tree/binary_tree_node_sum.py

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111

1212
from __future__ import annotations
1313

14+
from collections.abc import Iterator
15+
1416

1517
class Node:
1618
"""
@@ -33,24 +35,24 @@ class BinaryTreeNodeSum:
3335
12 8 0
3436
3537
>>> tree = Node(10)
36-
>>> BinaryTreeNodeSum(tree).node_sum()
38+
>>> sum(BinaryTreeNodeSum(tree))
3739
10
3840
3941
>>> tree.left = Node(5)
40-
>>> BinaryTreeNodeSum(tree).node_sum()
42+
>>> sum(BinaryTreeNodeSum(tree))
4143
15
4244
4345
>>> tree.right = Node(-3)
44-
>>> BinaryTreeNodeSum(tree).node_sum()
46+
>>> sum(BinaryTreeNodeSum(tree))
4547
12
4648
4749
>>> tree.left.left = Node(12)
48-
>>> BinaryTreeNodeSum(tree).node_sum()
50+
>>> sum(BinaryTreeNodeSum(tree))
4951
24
5052
5153
>>> tree.right.left = Node(8)
5254
>>> tree.right.right = Node(0)
53-
>>> BinaryTreeNodeSum(tree).node_sum()
55+
>>> sum(BinaryTreeNodeSum(tree))
5456
32
5557
"""
5658

@@ -64,8 +66,8 @@ def depth_first_search(self, node: Node | None) -> int:
6466
self.depth_first_search(node.left) + self.depth_first_search(node.right)
6567
)
6668

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)
6971

7072

7173
if __name__ == "__main__":

0 commit comments

Comments
 (0)