Coverage for blog/dsa/leetcode/bt/from_desc/__init__.py: 66%
119 statements
« prev ^ index » next coverage.py v7.6.12, created at 2025-02-20 16:23 +0000
1import json
2from typing import Any, Optional
4import pytest
5from typing_extensions import Self
8class TreeNode:
9 val: int
10 left: Self | None
11 right: Self | None
13 def __init__(
14 self,
15 val: int = 0,
16 left: Self | None = None,
17 right: Self | None = None,
18 ):
19 self.val = val
20 self.left = left
21 self.right = right
23 @classmethod
24 def fromDict(cls, values: dict[str, Any] | int | None) -> Self | None:
25 if values is None:
26 return None
27 elif isinstance(values, int):
28 return cls(values)
30 args = dict(val=values["value"])
31 if "left" in values:
32 args["left"] = cls.fromDict(values["left"])
33 if "right" in values:
34 args["right"] = cls.fromDict(values["right"])
35 return cls(**args)
37 def toDict(self) -> dict[str, Any]:
39 serial: dict[str, Any] = dict(value=self.val)
40 if self.left is not None:
41 serial["left"] = self.left.toDict()
42 if self.right is not None:
43 serial["right"] = self.right.toDict()
45 return serial
48# start snippet solution_initial
49class SolutionInitial:
50 def createBinaryTree(
51 self,
52 descriptions: list[list[int]],
53 ) -> Optional[TreeNode]:
55 # NOTE: Works because is tree of unique values.
56 nodes = {
57 color: (TreeNode(color), (color_parent, color, is_left))
58 for color_parent, color, is_left in descriptions
59 }
61 # NOTE: Linking step. Root should only appear as a parent, and thus
62 # was not constructed in the definition of nodes.
63 root = None
64 for node, (color_parent, _, is_left) in nodes.values():
65 if color_parent not in nodes:
66 if root is None:
67 root = TreeNode(color_parent)
68 node_parent = root
69 else:
70 node_parent, _ = nodes[color_parent]
72 if is_left:
73 node_parent.left = node
74 else:
75 node_parent.right = node
77 return root
78 # end snippet solution_initial
81# start snippet solution
82class SolutionGoodmem:
83 def createBinaryTree(
84 self,
85 descriptions: list[list[int]],
86 ) -> Optional[TreeNode]:
88 # NOTE: Works because is tree of unique values.
89 nodes = {}
90 nodes_no_parent = {}
91 for color_parent, color, is_left in descriptions:
92 node = TreeNode(color)
93 nodes[color] = node
95 # NOTE: Attach node to parent if found. Otherwise remember that it does
96 # not yet have a parent.
97 if color_parent not in nodes:
98 if color_parent not in nodes_no_parent:
99 nodes_no_parent[color_parent] = {is_left: node}
100 else:
101 nodes_no_parent[color_parent][is_left] = node
102 else:
103 node_parent = nodes[color_parent]
104 if is_left:
105 node_parent.left = node
106 else:
107 node_parent.right = node
109 # NOTE: If nodes_no_parent needs to link with this node, then do so
110 if node.val not in nodes_no_parent:
111 continue
113 for child_is_left, child in nodes_no_parent[node.val].items():
114 if child_is_left:
115 node.left = child
116 else:
117 node.right = child
119 nodes_no_parent.pop(node.val)
121 # There will only be one left.
122 for value, items in nodes_no_parent.items():
124 root = TreeNode(value)
125 if 0 in items:
126 root.right = items[0]
127 if 1 in items:
128 root.left = items[1]
130 return root
132 return None
133 # end snippet solution
136# start snippet solution_spoiled
137class Solution:
138 def createBinaryTree(
139 self,
140 descriptions: list[list[int]],
141 ) -> Optional[TreeNode]:
143 # NOTE: Works because is tree of unique values.
144 nodes = {}
145 children = set()
146 for desc in descriptions:
147 color_parent, color, is_left = desc
148 if color not in nodes:
149 node = TreeNode(color)
150 nodes[color] = node
151 else:
152 node = nodes[color]
154 # NOTE: Create parent node if not exists.
155 if color_parent not in nodes:
156 node_parent = TreeNode(color_parent)
157 nodes[color_parent] = node_parent
158 else:
159 node_parent = nodes[color_parent]
161 # NOTE: Assign.
162 if is_left:
163 node_parent.left = node
164 else:
165 node_parent.right = node
167 children.add(color)
169 # There will only be one left.
170 for value in nodes:
171 if value not in children:
172 return nodes[value]
174 return None
175 # end snippet solution_spoiled
178@pytest.fixture
179def solution():
180 return Solution()
183cases = (
184 (
185 [[20, 15, 1], [20, 17, 0], [50, 20, 1], [50, 80, 0], [80, 19, 1]],
186 treedict := {
187 "value": 50,
188 "left": {
189 "value": 20,
190 "left": {"value": 15},
191 "right": {"value": 17},
192 },
193 "right": {
194 "value": 80,
195 "left": {"value": 19},
196 },
197 },
198 ),
199 (
200 [[85, 74, 0], [38, 82, 0], [39, 70, 0], [82, 85, 0], [74, 13, 0], [13, 39, 0]],
201 {
202 "value": 38,
203 "right": {
204 "value": 82,
205 "right": {
206 "value": 85,
207 "right": {
208 "value": 74,
209 "right": {
210 "value": 13,
211 "right": {"value": 39, "right": {"value": 70}},
212 },
213 },
214 },
215 },
216 },
217 ),
218)
221def test_tree():
223 root = TreeNode.fromDict(treedict)
224 assert root is not None
225 assert root.val == 50
227 assert root.toDict() == treedict
230@pytest.mark.parametrize("question, answer", cases)
231def test_solution(
232 solution: Solution,
233 question: list[list[int]],
234 answer: dict[str, Any],
235):
237 got = solution.createBinaryTree(question)
238 answer_root = TreeNode.fromDict(answer)
239 assert answer_root is not None and got is not None
241 print("======================================")
242 print("answer")
243 print(json.dumps(answer_root.toDict(), indent=2))
245 print("======================================")
246 print("got")
247 print(json.dumps(got.toDict(), indent=2))
249 assert answer_root is not None
250 assert got.toDict() == answer_root.toDict()