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