Fixes; is_valid_tree >>> assert_tree

Author: nickolashkraus

src/merge_k_sorted_lists.py

@@ -65,6 +65,7 @@
 """
 
 import heapq
+from typing import Self
 
 from src.classes import ListNode
 
@@ -79,7 +80,7 @@ class HeapNode:
     def __init__(self, node: ListNode):
         self.node: ListNode = node
 
-    def __lt__(self, other: "HeapNode") -> bool:
+    def __lt__(self, other: Self) -> bool:
         return self.node.val < other.node.val
 
 

tests/test_construct_binary_tree_from_preorder_and_inorder_traversal.py

@@ -7,18 +7,18 @@
 from unittest import TestCase
 
 from src.construct_binary_tree_from_preorder_and_inorder_traversal import Solution
-from tests.utils import is_valid_tree
+from tests.utils import assert_tree
 
 
 class TestSolution(TestCase):
     def test_1(self):
         exp = [3, 9, 20, None, None, 15, 7]
-        is_valid_tree(Solution().buildTree(preorder=[3, 9, 20, 15, 7], inorder=[9, 3, 15, 20, 7]), 0, exp)
+        assert_tree(Solution().buildTree(preorder=[3, 9, 20, 15, 7], inorder=[9, 3, 15, 20, 7]), exp)
 
     def test_2(self):
         exp = [-1]
-        is_valid_tree(Solution().buildTree(preorder=[-1], inorder=[-1]), 0, exp)
+        assert_tree(Solution().buildTree(preorder=[-1], inorder=[-1]), exp)
 
     def test_3(self):
         exp = [1, None, 2]
-        is_valid_tree(Solution().buildTree(preorder=[1, 2], inorder=[1, 2]), 0, exp)
+        assert_tree(Solution().buildTree(preorder=[1, 2], inorder=[1, 2]), exp)

tests/utils.py

@@ -131,13 +131,17 @@ def create_bfs_list_from_binary_tree(*, root: TreeNode | None, values_only: bool
     return l
 
 
-def is_valid_tree(root: TreeNode | None, i: int, exp: list[T | None]) -> None:
+def assert_tree(root: TreeNode | None, exp: list[T | None]) -> None:
     """
-    Given a list of values in pre-order traversal, assert that the tree is
-    valid. Gaps in the tree are denoted by None.
+    Given a list of values in level-order traversal, assert that the tree
+    matches the expected values. Gaps in the tree should be denoted by None.
     """
-    if not root:
-        return None
-    assert root.val == exp[i]
-    is_valid_tree(root.left, 2 * i + 1, exp)
-    is_valid_tree(root.right, 2 * i + 2, exp)
+
+    def _assert_tree(root: TreeNode | None, i: int, exp: list[T | None]) -> None:
+        if not root:
+            return None
+        assert root.val == exp[i]
+        _assert_tree(root.left, 2 * i + 1, exp)
+        _assert_tree(root.right, 2 * i + 2, exp)
+
+    _assert_tree(root, 0, exp)