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 

3 

4import pytest 

5from typing_extensions import Self 

6 

7 

8class TreeNode: 

9 val: int 

10 left: Self | None 

11 right: Self | None 

12 

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 

22 

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) 

29 

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) 

36 

37 def toDict(self) -> dict[str, Any]: 

38 

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

44 

45 return serial 

46 

47 

48# start snippet solution_initial 

49class SolutionInitial: 

50 def createBinaryTree( 

51 self, 

52 descriptions: list[list[int]], 

53 ) -> Optional[TreeNode]: 

54 

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 } 

60 

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] 

71 

72 if is_left: 

73 node_parent.left = node 

74 else: 

75 node_parent.right = node 

76 

77 return root 

78 # end snippet solution_initial 

79 

80 

81# start snippet solution 

82class SolutionGoodmem: 

83 def createBinaryTree( 

84 self, 

85 descriptions: list[list[int]], 

86 ) -> Optional[TreeNode]: 

87 

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 

94 

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 

108 

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 

112 

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 

118 

119 nodes_no_parent.pop(node.val) 

120 

121 # There will only be one left. 

122 for value, items in nodes_no_parent.items(): 

123 

124 root = TreeNode(value) 

125 if 0 in items: 

126 root.right = items[0] 

127 if 1 in items: 

128 root.left = items[1] 

129 

130 return root 

131 

132 return None 

133 # end snippet solution 

134 

135 

136# start snippet solution_spoiled 

137class Solution: 

138 def createBinaryTree( 

139 self, 

140 descriptions: list[list[int]], 

141 ) -> Optional[TreeNode]: 

142 

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] 

153 

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] 

160 

161 # NOTE: Assign. 

162 if is_left: 

163 node_parent.left = node 

164 else: 

165 node_parent.right = node 

166 

167 children.add(color) 

168 

169 # There will only be one left. 

170 for value in nodes: 

171 if value not in children: 

172 return nodes[value] 

173 

174 return None 

175 # end snippet solution_spoiled 

176 

177 

178@pytest.fixture 

179def solution(): 

180 return Solution() 

181 

182 

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) 

219 

220 

221def test_tree(): 

222 

223 root = TreeNode.fromDict(treedict) 

224 assert root is not None 

225 assert root.val == 50 

226 

227 assert root.toDict() == treedict 

228 

229 

230@pytest.mark.parametrize("question, answer", cases) 

231def test_solution( 

232 solution: Solution, 

233 question: list[list[int]], 

234 answer: dict[str, Any], 

235): 

236 

237 got = solution.createBinaryTree(question) 

238 answer_root = TreeNode.fromDict(answer) 

239 assert answer_root is not None and got is not None 

240 

241 print("======================================") 

242 print("answer") 

243 print(json.dumps(answer_root.toDict(), indent=2)) 

244 

245 print("======================================") 

246 print("got") 

247 print(json.dumps(got.toDict(), indent=2)) 

248 

249 assert answer_root is not None 

250 assert got.toDict() == answer_root.toDict()