diff --git a/src/binary_tree/dfs/count_good_nodes_in_binary_tree/__init__.py b/src/binary_tree/dfs/count_good_nodes_in_binary_tree/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/binary_tree/dfs/count_good_nodes_in_binary_tree/solution.py b/src/binary_tree/dfs/count_good_nodes_in_binary_tree/solution.py new file mode 100644 index 0000000..e47dc37 --- /dev/null +++ b/src/binary_tree/dfs/count_good_nodes_in_binary_tree/solution.py @@ -0,0 +1,23 @@ +from structures import TreeNode + + +class Solution: + def good_w_max(self, cur_max, root: TreeNode) -> int: + if root: + counter = 0 + if root.val >= cur_max: + counter += 1 + cur_max = max(cur_max, root.val) + return ( + counter + + self.good_w_max(cur_max, root.left) + + self.good_w_max(cur_max, root.right) + ) + return 0 + + def goodNodes(self, root: TreeNode) -> int: + return ( + 1 + + self.good_w_max(root.val, root.left) + + self.good_w_max(root.val, root.right) + ) diff --git a/tests/test_count-good-nodes-in-binary-tree.py b/tests/test_count-good-nodes-in-binary-tree.py new file mode 100644 index 0000000..6495831 --- /dev/null +++ b/tests/test_count-good-nodes-in-binary-tree.py @@ -0,0 +1,13 @@ +import pytest +from src.binary_tree.dfs.count_good_nodes_in_binary_tree.solution import Solution +from test_utils import TreeBuilder + + +@pytest.mark.parametrize( + "root, expected", + [([3, 1, 4, 3, None, 1, 5], 4), ([3, 3, None, 4, 2], 3), ([1], 1)], +) +def test_max_depth(root, expected): + root = TreeBuilder.from_list(root) + solution = Solution() + assert solution.goodNodes(root) == expected