def search_range(self, lower, upper, return_key=False):
"""
Returns a list of all values whose keys are in the range [lower, upper] inclusive
"""
if lower > upper:
return []
first_gte = self.search_first_gte(lower)
res = []
if first_gte == None:
return res
node, pos = first_gte
while node:
for i in range(pos, len(node.keys)):
if node.keys[i] > upper:
# current and all other leaf nodes on the road are greater than upper bound and not part of res
# so we can just return res
return res
if return_key:
res.append(node.keys[i])
else:
res.append(node.pointers[i])
# move to the immediate right neighbour
if node.pointers[-1] == None:
return res
node = node.pointers[-1]
Tracker.add_to_set("leaf", node)
pos = 0
# this return is needed if the res includes the rightmost leaf node
return res
def search_first_gte(self, key):
"""
A utility function used by search_range to return the first leaf node >= key
If found, return the leaf node containing the key and the index of the key in the node
If not found, i.e. key is smaller than all keys, return None
"""
if self.leaf:
Tracker.add_to_set("leaf", self)
for i in range(len(self.keys)):
if self.keys[i] >= key:
return self, i
if self.pointers[-1] == None:
# this is true if self is the rightmost leaf node
return None
# if leaf node is not rightmost, we know the first key of the immediate right neightbour will satisfy condition
# because self.pointers[-1].keys[0] >= some LB > key
return self.pointers[-1], 0
else:
Tracker.add_to_set("non-leaf", self)
# find the subtree to recursively call on
for i in range(len(self.keys)):
if key < self.keys[i]:
return self.pointers[i].search_first_gte(key)
return self.pointers[-1].search_first_gte(key)