Skip to content

Commit c8024c1

Browse files
Add Bit-Merging Tree (BMTree) data structure
1 parent adccf0c commit c8024c1

1 file changed

Lines changed: 298 additions & 0 deletions

File tree

Lines changed: 298 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,298 @@
1+
"""
2+
Bit-Merging Tree (BMTree) Implementation.
3+
4+
A Bit-Merging Tree (BMTree) indexes records that have three attributes (such as EV
5+
voltage, temperature, and power) by merging (interleaving) the bits of the attributes
6+
into a single 30-bit Morton (Z-order) key. Records with similar attributes produce
7+
similar keys, grouping them close together in the hierarchical tree structure. Keys
8+
whose distance to their nearest neighbor exceeds a threshold (outliers) are stored
9+
separately so they do not distort the primary tree index.
10+
11+
References:
12+
https://en.wikipedia.org/wiki/Z-order_curve
13+
https://en.wikipedia.org/wiki/Binary_search_tree
14+
"""
15+
16+
from __future__ import annotations
17+
18+
from bisect import bisect_left
19+
from dataclasses import dataclass, field
20+
21+
DIMENSIONS = 3
22+
BITS_PER_DIMENSION = 10
23+
MAX_ATTRIBUTE_VALUE = 2**BITS_PER_DIMENSION - 1 # 1023
24+
25+
26+
@dataclass
27+
class BMNode:
28+
"""
29+
A single node in the BMTree holding one Morton key.
30+
31+
>>> node = BMNode(7)
32+
>>> node.key, node.left, node.right
33+
(7, None, None)
34+
"""
35+
36+
key: int
37+
left: BMNode | None = None
38+
right: BMNode | None = None
39+
40+
41+
@dataclass
42+
class BMTree:
43+
"""
44+
A BMTree container holding the main clustered tree and separately stored outliers.
45+
46+
>>> tree = BMTree()
47+
>>> tree.root is None, tree.outliers
48+
(True, [])
49+
"""
50+
51+
root: BMNode | None = None
52+
outliers: list[int] = field(default_factory=list)
53+
54+
55+
def encode_morton_key(point: tuple[int, ...]) -> int:
56+
"""
57+
Merge the bits of the three attributes into one 30-bit Morton key.
58+
59+
>>> encode_morton_key((1, 0, 0))
60+
4
61+
>>> encode_morton_key((0, 1, 0))
62+
2
63+
>>> encode_morton_key((0, 0, 1))
64+
1
65+
>>> encode_morton_key((0, 0, 2))
66+
8
67+
>>> encode_morton_key((1023, 1023, 1023))
68+
1073741823
69+
>>> encode_morton_key((1024, 0, 0))
70+
Traceback (most recent call last):
71+
...
72+
ValueError: attribute values must be integers between 0 and 1023
73+
>>> encode_morton_key((-1, 0, 0))
74+
Traceback (most recent call last):
75+
...
76+
ValueError: attribute values must be integers between 0 and 1023
77+
>>> encode_morton_key((1, 2))
78+
Traceback (most recent call last):
79+
...
80+
ValueError: a point must have exactly 3 attributes
81+
"""
82+
if len(point) != DIMENSIONS:
83+
msg = f"a point must have exactly {DIMENSIONS} attributes"
84+
raise ValueError(msg)
85+
if any(
86+
not isinstance(value, int) or not 0 <= value <= MAX_ATTRIBUTE_VALUE
87+
for value in point
88+
):
89+
msg = f"attribute values must be integers between 0 and {MAX_ATTRIBUTE_VALUE}"
90+
raise ValueError(msg)
91+
92+
key = 0
93+
for bit_position in range(BITS_PER_DIMENSION - 1, -1, -1):
94+
for value in point:
95+
key = (key << 1) | ((value >> bit_position) & 1)
96+
return key
97+
98+
99+
def decode_morton_key(key: int) -> tuple[int, int, int]:
100+
"""
101+
Split a 30-bit Morton key back into its three constituent attributes.
102+
103+
>>> decode_morton_key(4)
104+
(1, 0, 0)
105+
>>> decode_morton_key(encode_morton_key((365, 48, 550)))
106+
(365, 48, 550)
107+
>>> decode_morton_key(2**30)
108+
Traceback (most recent call last):
109+
...
110+
ValueError: key must be between 0 and 1073741823
111+
"""
112+
max_key = 2 ** (BITS_PER_DIMENSION * DIMENSIONS) - 1
113+
if not 0 <= key <= max_key:
114+
msg = f"key must be between 0 and {max_key}"
115+
raise ValueError(msg)
116+
117+
values = [0] * DIMENSIONS
118+
for index in range(BITS_PER_DIMENSION * DIMENSIONS):
119+
bit = (key >> index) & 1
120+
dimension = DIMENSIONS - 1 - (index % DIMENSIONS)
121+
values[dimension] |= bit << (index // DIMENSIONS)
122+
return values[0], values[1], values[2]
123+
124+
125+
def split_outliers(
126+
sorted_keys: list[int], outlier_gap: int
127+
) -> tuple[list[int], list[int]]:
128+
"""
129+
Separate keys whose distance to their nearest neighboring key exceeds outlier_gap.
130+
Expects sorted, unique keys. Returns (clustered_keys, outlier_keys).
131+
132+
>>> split_outliers([1, 2, 3, 100], 5)
133+
([1, 2, 3], [100])
134+
>>> split_outliers([5], 5)
135+
([5], [])
136+
>>> split_outliers([], 5)
137+
([], [])
138+
>>> split_outliers([1, 2, 3], 5)
139+
([1, 2, 3], [])
140+
"""
141+
clustered: list[int] = []
142+
outliers: list[int] = []
143+
for index, key in enumerate(sorted_keys):
144+
gaps = []
145+
if index > 0:
146+
gaps.append(key - sorted_keys[index - 1])
147+
if index < len(sorted_keys) - 1:
148+
gaps.append(sorted_keys[index + 1] - key)
149+
if gaps and min(gaps) > outlier_gap:
150+
outliers.append(key)
151+
else:
152+
clustered.append(key)
153+
return clustered, outliers
154+
155+
156+
def _build_balanced(sorted_keys: list[int], low: int, high: int) -> BMNode | None:
157+
"""
158+
Recursively build a balanced binary search tree from sorted_keys[low:high + 1].
159+
160+
>>> _build_balanced([], 0, -1) is None
161+
True
162+
>>> node = _build_balanced([1, 2, 3], 0, 2)
163+
>>> node is not None and node.key == 2
164+
True
165+
"""
166+
if low > high:
167+
return None
168+
middle = (low + high) // 2
169+
return BMNode(
170+
sorted_keys[middle],
171+
_build_balanced(sorted_keys, low, middle - 1),
172+
_build_balanced(sorted_keys, middle + 1, high),
173+
)
174+
175+
176+
def build_bm_tree(
177+
points: list[tuple[int, int, int]], outlier_gap: int | None = None
178+
) -> BMTree:
179+
"""
180+
Build a BMTree from a list of three-attribute points.
181+
182+
If outlier_gap is provided, keys farther than outlier_gap from their nearest
183+
neighbor are isolated in the outliers list instead of the primary tree.
184+
185+
>>> tree = build_bm_tree([(1, 1, 1), (1, 1, 2), (1, 2, 1)])
186+
>>> tree.root is not None and tree.root.key == 14
187+
True
188+
>>> tree.outliers
189+
[]
190+
>>> points = [(1, 1, 1), (1, 1, 2), (1, 2, 1), (1000, 1000, 1000)]
191+
>>> tree = build_bm_tree(points, outlier_gap=100)
192+
>>> tree.root is not None and tree.root.key == 14
193+
True
194+
>>> tree.outliers == [encode_morton_key((1000, 1000, 1000))]
195+
True
196+
>>> build_bm_tree([]).root is None
197+
True
198+
>>> build_bm_tree([(1, 1, 1)], outlier_gap=-1)
199+
Traceback (most recent call last):
200+
...
201+
ValueError: outlier_gap must not be negative
202+
>>> build_bm_tree([(1, 1, 5000)])
203+
Traceback (most recent call last):
204+
...
205+
ValueError: attribute values must be integers between 0 and 1023
206+
"""
207+
if outlier_gap is not None and outlier_gap < 0:
208+
msg = "outlier_gap must not be negative"
209+
raise ValueError(msg)
210+
211+
keys = sorted({encode_morton_key(point) for point in points})
212+
if outlier_gap is None:
213+
clustered, outliers = keys, []
214+
else:
215+
clustered, outliers = split_outliers(keys, outlier_gap)
216+
return BMTree(_build_balanced(clustered, 0, len(clustered) - 1), outliers)
217+
218+
219+
def _tree_contains(root: BMNode | None, key: int) -> bool:
220+
"""
221+
Return True if key is stored in the tree rooted at root, False otherwise.
222+
223+
>>> _tree_contains(None, 1)
224+
False
225+
>>> _tree_contains(BMNode(5, BMNode(2), BMNode(9)), 9)
226+
True
227+
>>> _tree_contains(BMNode(5, BMNode(2), BMNode(9)), 7)
228+
False
229+
"""
230+
current = root
231+
while current is not None:
232+
if key == current.key:
233+
return True
234+
current = current.left if key < current.key else current.right
235+
return False
236+
237+
238+
def is_outlier(tree: BMTree, point: tuple[int, int, int]) -> bool:
239+
"""
240+
Return True if the point is stored in the outlier list of the BMTree.
241+
242+
>>> points = [(1, 1, 1), (1, 1, 2), (1, 2, 1), (1000, 1000, 1000)]
243+
>>> tree = build_bm_tree(points, outlier_gap=100)
244+
>>> is_outlier(tree, (1000, 1000, 1000))
245+
True
246+
>>> is_outlier(tree, (1, 1, 1))
247+
False
248+
"""
249+
key = encode_morton_key(point)
250+
index = bisect_left(tree.outliers, key)
251+
return index < len(tree.outliers) and tree.outliers[index] == key
252+
253+
254+
def search(tree: BMTree, point: tuple[int, int, int]) -> bool:
255+
"""
256+
Search whether a given three-attribute point exists in the BMTree
257+
(either within the primary clustered tree or the outlier storage).
258+
259+
>>> points = [(365, 48, 550), (366, 48, 551), (900, 10, 20)]
260+
>>> tree = build_bm_tree(points)
261+
>>> search(tree, (365, 48, 550))
262+
True
263+
>>> search(tree, (900, 10, 20))
264+
True
265+
>>> search(tree, (365, 48, 551))
266+
False
267+
>>> search(build_bm_tree([]), (1, 2, 3))
268+
False
269+
>>> search(tree, (1, 2))
270+
Traceback (most recent call last):
271+
...
272+
ValueError: a point must have exactly 3 attributes
273+
>>> tree = build_bm_tree([(1, 1, 1), (1, 1, 2), (1000, 1000, 1000)], 100)
274+
>>> search(tree, (1000, 1000, 1000))
275+
True
276+
"""
277+
key = encode_morton_key(point)
278+
return _tree_contains(tree.root, key) or is_outlier(tree, point)
279+
280+
281+
def inorder_traversal(root: BMNode | None) -> list[int]:
282+
"""
283+
Return the Morton keys stored in the tree in ascending in-order sequence.
284+
285+
>>> inorder_traversal(None)
286+
[]
287+
>>> inorder_traversal(BMNode(5, BMNode(2), BMNode(9)))
288+
[2, 5, 9]
289+
"""
290+
if root is None:
291+
return []
292+
return [*inorder_traversal(root.left), root.key, *inorder_traversal(root.right)]
293+
294+
295+
if __name__ == "__main__":
296+
import doctest
297+
298+
doctest.testmod()

0 commit comments

Comments
 (0)