Skip to content

Commit 8a2ec0d

Browse files
committed
Add Radix Binary Search with documentation and doctests
1 parent 6e03a48 commit 8a2ec0d

1 file changed

Lines changed: 169 additions & 0 deletions

File tree

‎searches/radix_binary_search.py‎

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
"""Radix Binary Search: a prefix directory followed by bounded binary search.
2+
3+
Subtract the minimum value to obtain nonnegative keys. Their leading bits
4+
identify buckets containing contiguous ranges of the sorted input. Store each
5+
bucket's start, then use two adjacent starts as the bounds for binary search.
6+
7+
For example, [-8, -8, -1, 0, 7, 8] with bits=2 has minimum -8 and span 16.
8+
The normalized keys are [0, 0, 7, 8, 15, 16]. Using their common five-bit
9+
width, discard three low bits to get bucket numbers [0, 0, 0, 1, 1, 2].
10+
The directory is [0, 3, 5, 6], so bucket 1 covers indices [3, 5): [0, 7].
11+
12+
Background: Kipf et al. (2019), "SOSD: A Benchmark for Learned Indexes",
13+
Section 2: RadixBinarySearch (RBS).
14+
https://arxiv.org/abs/1911.13014
15+
"""
16+
17+
from collections.abc import Sequence
18+
from itertools import pairwise
19+
20+
21+
class RadixBinarySearch:
22+
"""Prepare a reusable index over integers sorted in nondecreasing order.
23+
24+
Keep a tuple snapshot so later changes to the caller's list cannot make
25+
the directory stale. Results are indices in that snapshot. Duplicates
26+
return their first index; missing targets return None.
27+
28+
``bits`` controls the maximum number of buckets: 2**bits. The limit of
29+
12 bounds the directory to 4096 buckets; zero uses a single bucket.
30+
This cap is a memory-budget choice, not a requirement of the algorithm.
31+
32+
For n values and B buckets, preparation takes O(n + B) time and storage,
33+
including the snapshot. Searching a bucket with m values takes
34+
O(1 + log(m + 1)) time, O(log(n + 1)) worst case, and O(1) extra space.
35+
These bounds assume constant-time integer arithmetic and comparisons.
36+
Skewed data can put most values into one bucket. Reuse the index across
37+
queries to amortize preparation; faster lookup is not guaranteed.
38+
"""
39+
40+
def __init__(self, values: Sequence[int], *, bits: int = 12) -> None:
41+
"""Validate the input, copy it, and build bucket boundaries.
42+
43+
>>> RadixBinarySearch([2, 1])
44+
Traceback (most recent call last):
45+
...
46+
ValueError: values must be sorted in nondecreasing order
47+
>>> RadixBinarySearch([1, 2.5])
48+
Traceback (most recent call last):
49+
...
50+
TypeError: values must contain integers, excluding bool
51+
>>> RadixBinarySearch([], bits=13)
52+
Traceback (most recent call last):
53+
...
54+
ValueError: bits must be between 0 and 12
55+
>>> RadixBinarySearch([], bits=True)
56+
Traceback (most recent call last):
57+
...
58+
TypeError: bits must be an integer, excluding bool
59+
"""
60+
# bool is an int subclass in Python, but is not a meaningful bit count.
61+
# Validate even for empty input so the parameter contract is consistent.
62+
if isinstance(bits, bool) or not isinstance(bits, int):
63+
raise TypeError("bits must be an integer, excluding bool")
64+
if not 0 <= bits <= 12:
65+
raise ValueError("bits must be between 0 and 12")
66+
67+
# Retain input positions: sorting here would change the returned indices.
68+
# A tuple keeps this snapshot stable if the caller later edits their list.
69+
self._values = tuple(values)
70+
if any(
71+
isinstance(value, bool) or not isinstance(value, int)
72+
for value in self._values
73+
):
74+
raise TypeError("values must contain integers, excluding bool")
75+
if any(left > right for left, right in pairwise(self._values)):
76+
raise ValueError("values must be sorted in nondecreasing order")
77+
78+
# Empty input uses a harmless zero-width range and offsets [0, 0].
79+
# search() checks for emptiness before accessing either endpoint.
80+
self._minimum = self._values[0] if self._values else 0
81+
maximum = self._values[-1] if self._values else 0
82+
span = maximum - self._minimum
83+
84+
# Subtracting the minimum preserves order and makes every key nonnegative.
85+
# All keys share the width implied by span; smaller keys have leading zeros.
86+
# Right-shifting by (width - bits) keeps at most `bits` prefix bits.
87+
# A zero shift keeps all bits. With bits=0 every key maps to bucket 0.
88+
self._shift = max(0, span.bit_length() - bits)
89+
# Bucket numbers start at zero, so include the maximum key's bucket.
90+
bucket_count = (span >> self._shift) + 1
91+
92+
# The cursor only moves forward: total work is O(n + B).
93+
# Equal adjacent offsets describe an empty bucket. The extra offset
94+
# at the end is len(values), closing the final half-open interval.
95+
self._offsets: list[int] = []
96+
position = 0
97+
for bucket in range(bucket_count + 1):
98+
# A left shift converts a prefix back to its minimum normalized key.
99+
# Adding the minimum restores the corresponding input value.
100+
boundary = self._minimum + (bucket << self._shift)
101+
# Stop BEFORE equal values: every duplicate stays in the same bucket.
102+
# offsets[bucket] is the first index with value >= boundary.
103+
while position < len(self._values) and self._values[position] < boundary:
104+
position += 1
105+
self._offsets.append(position)
106+
107+
def search(self, target: int) -> int | None:
108+
"""Return the first matching index, or None if target is absent.
109+
110+
>>> index = RadixBinarySearch([-8, -8, -1, 0, 7, 8], bits=2)
111+
>>> [index.search(target) for target in (-8, -1, 0, 7, 8)]
112+
[0, 2, 3, 4, 5]
113+
>>> [index.search(target) for target in (-9, 3, 9)]
114+
[None, None, None]
115+
>>> RadixBinarySearch([], bits=0).search(1) is None
116+
True
117+
>>> RadixBinarySearch([5, 5, 5], bits=0).search(5)
118+
0
119+
>>> RadixBinarySearch([-4, 0, 0, 9], bits=0).search(0)
120+
1
121+
>>> RadixBinarySearch([0, 100], bits=4).search(50) is None
122+
True
123+
>>> RadixBinarySearch([-10**100, 0, 10**100]).search(10**100)
124+
2
125+
>>> values = [1, 2, 3]
126+
>>> index = RadixBinarySearch(values)
127+
>>> values[0] = 99
128+
>>> index.search(1)
129+
0
130+
>>> index.search(True)
131+
Traceback (most recent call last):
132+
...
133+
TypeError: target must be an integer, excluding bool
134+
"""
135+
if isinstance(target, bool) or not isinstance(target, int):
136+
raise TypeError("target must be an integer, excluding bool")
137+
# Check the range before calculating a bucket: a smaller target could
138+
# produce a negative Python index, and a larger one could exceed the table.
139+
if not self._values or target < self._minimum or target > self._values[-1]:
140+
return None
141+
142+
bucket = (target - self._minimum) >> self._shift
143+
left = self._offsets[bucket]
144+
stop = self._offsets[bucket + 1]
145+
# Search the half-open interval [left, right), without copying a slice.
146+
# Preserve stop because right shrinks during the lower-bound search.
147+
right = stop
148+
149+
# Discard values strictly below target. On equality, retain the left half
150+
# so an earlier duplicate can still be found. The result is the first
151+
# index in this bucket whose value is >= target, or stop if none exists.
152+
while left < right:
153+
middle = (left + right) // 2
154+
if self._values[middle] < target:
155+
left = middle + 1
156+
else:
157+
right = middle
158+
159+
# A lower bound can point at a larger value or at the end of the bucket.
160+
# Verify equality; the bounds check also handles empty buckets safely.
161+
if left < stop and self._values[left] == target:
162+
return left
163+
return None
164+
165+
166+
if __name__ == "__main__":
167+
import doctest
168+
169+
doctest.testmod()

0 commit comments

Comments
 (0)