Coverage for blog/dsa/bst/__init__.py: 93%
151 statements
« prev ^ index » next coverage.py v7.6.12, created at 2025-02-20 16:23 +0000
« prev ^ index » next coverage.py v7.6.12, created at 2025-02-20 16:23 +0000
1import json
2import random
3import secrets
4from typing import Any, Generator, Iterable
6from typing_extensions import Iterator, Self
8from dsa import queue
11class Node:
12 left: Self | None
13 right: Self | None
14 value: int
15 identity: str
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)
30 @classmethod
31 def mk(cls, max_size: int, *, value_max: int = 1000, value_min: int = 0) -> Self:
33 root = cls(random.randint(value_min, value_max))
34 for _ in range(max_size):
35 root.add(random.randint(value_min, value_max))
37 return root
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__()
46 # ----------------------------------------------------------------------- #
47 # NOTE: For play and tests.
49 def values(self) -> list[int]:
50 return list(item.value for item in self)
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()
59 return out
61 @classmethod
62 def from_dict(cls, raw: dict[str, Any]):
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"])
70 return Node(**out)
72 def dump_json(self, **kwargs):
73 return json.dumps(self.to_dict(), **kwargs)
75 @classmethod
76 def load_json(cls, raw: str):
77 return Node.from_dict(json.loads(raw))
79 # ----------------------------------------------------------------------- #
80 # NOTE: Methods that might show up in a test/interview.
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
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
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
106 return self
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."""
112 if (node_parent := self.find(value, parent=True)) is None:
113 return None
115 node = node_parent.find(value)
116 if node is None:
117 raise ValueError("Found parent of node but no node.")
119 if node_parent.left == node:
120 node_parent.left = None
121 else:
122 node_parent.right = None
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
131 return node
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."""
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
152 def _size(self, count: int = 0):
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)
160 return count
162 def size(self) -> int:
163 """Compute the size of the tree."""
164 return self._size()
166 def min(self) -> Self:
167 """Minimum should be the leftmost node on the tree.
169 Obviously, we could just use ``__iter__`` and take the minimum of
170 some sequence, but this is certainly less efficient.
171 """
173 if self.left:
174 return self.left.min()
175 else:
176 return self
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
184 for node in self:
185 if not (diff := abs(node.value - value)):
186 best = node
187 break
189 if diff_best is None or diff_best > diff:
190 best = node
191 diff_best = diff
193 return best # type: ignore[return-value]
195 def _depth(self, depth_last: int = 0, depth_best: int = 0) -> int:
197 depth_last += 1
198 if depth_last > depth_best:
199 depth_best = depth_last
201 if self.left is not None:
202 depth_best = self.left._depth(depth_last, depth_best)
204 if self.right is not None:
205 depth_best = self.right._depth(depth_last, depth_best)
207 return depth_best
209 def depth(self):
210 return self._depth(0, 0)
212 @classmethod
213 def _iter_layers(cls, prev_layer: Iterable[Self]) -> queue.Queue[Self]:
214 """Provided the previous layer of nodes, get the next."""
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)
223 return queue_next
225 def iter_layers(self) -> Generator[tuple[int, tuple[Self, ...]], None, None]:
226 queue_current = queue.Queue[Self]()
227 queue_current.enqueue(self)
229 depth = -1
230 while queue_current:
231 # NOTE: Do not spent the queue here!
232 yield (depth := depth + 1, tuple(queue_current.values()))
234 # NOTE: Because the queue is spent here.
235 queue_current = self._iter_layers(queue_current)
237 def iter_bredth(self) -> Generator[Self, None, None]:
238 for _, layer in self.iter_layers():
239 yield from layer