Coverage for blog/dsa/bst/__init__.py: 93%

1import json 

2import random 

3import secrets 

4from typing import Any, Generator, Iterable 

5 

6from typing_extensions import Iterator, Self 

7 

8from dsa import queue 

9 

10 

11class Node: 

12 left: Self | None 

13 right: Self | None 

14 value: int 

15 identity: str 

16 

17 def __init__( 

18 self, 

19 value: int, 

20 *, 

21 left: Self | None = None, 

22 right: Self | None = None, 

23 identity: str | None = None, 

24 ): 

25 self.left = left 

26 self.right = right 

27 self.value = value 

28 self.identity = identity if identity is not None else secrets.token_hex(8) 

29 

30 @classmethod 

31 def mk(cls, max_size: int, *, value_max: int = 1000, value_min: int = 0) -> Self: 

32 

33 root = cls(random.randint(value_min, value_max)) 

34 for _ in range(max_size): 

35 root.add(random.randint(value_min, value_max)) 

36 

37 return root 

38 

39 def __iter__(self) -> Iterator[Self]: 

40 yield self 

41 if self.left is not None: 

42 yield from self.left.__iter__() 

43 if self.right is not None: 

44 yield from self.right.__iter__() 

45 

46 # ----------------------------------------------------------------------- # 

47 # NOTE: For play and tests. 

48 

49 def values(self) -> list[int]: 

50 return list(item.value for item in self) 

51 

52 def to_dict(self) -> dict[str, Any]: 

53 out: dict[str, Any] = {"value": self.value} 

54 if self.right is not None: 

55 out["right"] = self.right.to_dict() 

56 if self.left is not None: 

57 out["left"] = self.left.to_dict() 

58 

59 return out 

60 

61 @classmethod 

62 def from_dict(cls, raw: dict[str, Any]): 

63 

64 out: dict[str, Any] = {"value": raw["value"]} 

65 if "left" in raw: 

66 out["left"] = Node.from_dict(raw["left"]) 

67 if "right" in raw: 

68 out["right"] = Node.from_dict(raw["right"]) 

69 

70 return Node(**out) 

71 

72 def dump_json(self, **kwargs): 

73 return json.dumps(self.to_dict(), **kwargs) 

74 

75 @classmethod 

76 def load_json(cls, raw: str): 

77 return Node.from_dict(json.loads(raw)) 

78 

79 # ----------------------------------------------------------------------- # 

80 # NOTE: Methods that might show up in a test/interview. 

81 

82 def add(self, value: int) -> Self: 

83 """Add a node.""" 

84 if value < self.value: 

85 if self.left is None: 

86 self.left = self.__class__(value) 

87 return self.left 

88 

89 return self.left.add(value) 

90 elif value > self.value: 

91 if self.right is None: 

92 self.right = self.__class__(value) 

93 return self.right 

94 return self.right.add(value) 

95 else: 

96 return self 

97 

98 def check(self) -> Self: 

99 """Is this tree a binary search tree?""" 

100 for node in self: 

101 if node.left is not None: 

102 assert node.left.value < node.value 

103 if node.right is not None: 

104 assert node.right.value > node.value 

105 

106 return self 

107 

108 # NOTE: This is done from intuition, but it would appear to work. 

109 def pop(self, value: int) -> Self | None: 

110 """Remove a value from the tree.""" 

111 

112 if (node_parent := self.find(value, parent=True)) is None: 

113 return None 

114 

115 node = node_parent.find(value) 

116 if node is None: 

117 raise ValueError("Found parent of node but no node.") 

118 

119 if node_parent.left == node: 

120 node_parent.left = None 

121 else: 

122 node_parent.right = None 

123 

124 if (left := node.left) is not None: 

125 left_new = self.add(left.value) 

126 left_new.right, left_new.left = left.right, left.left 

127 if (right := node.right) is not None: 

128 right_new = self.add(right.value) 

129 right_new.right, right_new.left = right.right, right.left 

130 

131 return node 

132 

133 def find( 

134 self, value: int, *, parent: bool = False, _prev: Self | None = None 

135 ) -> Self | None: 

136 """Find a value in the tree. Use ``parent`` to return the parent node.""" 

137 

138 if value < self.value: 

139 if self.left is None: 

140 return None 

141 return self.left.find(value, _prev=self, parent=parent) 

142 elif value > self.value: 

143 if self.right is None: 

144 return None 

145 return self.right.find(value, _prev=self, parent=parent) 

146 else: 

147 if parent: 

148 return _prev 

149 else: 

150 return self 

151 

152 def _size(self, count: int = 0): 

153 

154 count += 1 # count self 

155 if self.left is not None: 

156 count = self.left._size(count) 

157 if self.right is not None: 

158 count = self.right._size(count) 

159 

160 return count 

161 

162 def size(self) -> int: 

163 """Compute the size of the tree.""" 

164 return self._size() 

165 

166 def min(self) -> Self: 

167 """Minimum should be the leftmost node on the tree. 

168 

169 Obviously, we could just use ``__iter__`` and take the minimum of 

170 some sequence, but this is certainly less efficient. 

171 """ 

172 

173 if self.left: 

174 return self.left.min() 

175 else: 

176 return self 

177 

178 # Is there a more clever way to do this? This is basically 

179 # https://medium.com/journey-to-becoming-an-algoat/closest-value-in-a-bst-with-recursion-16bf90ad3bc2 

180 def approximate(self, value: int) -> Self: 

181 best: Node = self 

182 diff_best = None 

183 

184 for node in self: 

185 if not (diff := abs(node.value - value)): 

186 best = node 

187 break 

188 

189 if diff_best is None or diff_best > diff: 

190 best = node 

191 diff_best = diff 

192 

193 return best # type: ignore[return-value] 

194 

195 def _depth(self, depth_last: int = 0, depth_best: int = 0) -> int: 

196 

197 depth_last += 1 

198 if depth_last > depth_best: 

199 depth_best = depth_last 

200 

201 if self.left is not None: 

202 depth_best = self.left._depth(depth_last, depth_best) 

203 

204 if self.right is not None: 

205 depth_best = self.right._depth(depth_last, depth_best) 

206 

207 return depth_best 

208 

209 def depth(self): 

210 return self._depth(0, 0) 

211 

212 @classmethod 

213 def _iter_layers(cls, prev_layer: Iterable[Self]) -> queue.Queue[Self]: 

214 """Provided the previous layer of nodes, get the next.""" 

215 

216 queue_next = queue.Queue[Self]() 

217 for node in prev_layer: 

218 if node.left is not None: 

219 queue_next.enqueue(node.left) 

220 if node.right is not None: 

221 queue_next.enqueue(node.right) 

222 

223 return queue_next 

224 

225 def iter_layers(self) -> Generator[tuple[int, tuple[Self, ...]], None, None]: 

226 queue_current = queue.Queue[Self]() 

227 queue_current.enqueue(self) 

228 

229 depth = -1 

230 while queue_current: 

231 # NOTE: Do not spent the queue here! 

232 yield (depth := depth + 1, tuple(queue_current.values())) 

233 

234 # NOTE: Because the queue is spent here. 

235 queue_current = self._iter_layers(queue_current) 

236 

237 def iter_bredth(self) -> Generator[Self, None, None]: 

238 for _, layer in self.iter_layers(): 

239 yield from layer