Kth Smallest, Floor, Ceiling and Range Queries
32 min · Intermediate
Need help? Ask a tutor on WhatsApp (opens in a new tab)Because inorder gives sorted keys, a BST answers order questions: the third smallest, the nearest key to a price, all keys between two limits. The aim is to look at as few nodes as possible.
Kth smallest
The kth smallest key is the kth key of the inorder walk. Since the walk is a generator, islice can skip k - 1 keys and take the next one, and the rest of the tree is never visited.
from itertools import islicedef kth_smallest(root, k): return next(islice(inorder(root), k - 1, None))Example 1 · The three smallest keys
class Node: def __init__(self, key): self.key = key self.left = None self.right = None
from collections import deque
def build(tokens): if not tokens or tokens[0] == "#": return None root = Node(int(tokens[0])) queue = deque([root]) rest = iter(tokens[1:]) while queue: node = queue.popleft() for side in ("left", "right"): token = next(rest, "#") if token != "#": child = Node(int(token)) setattr(node, side, child) queue.append(child) return root
def inorder(node): if node is not None: yield from inorder(node.left) yield node.key yield from inorder(node.right)
from itertools import islicedef kth_smallest(root, k): return next(islice(inorder(root), k - 1, None))
root = build('8 3 10 1 6 # 14 # # 4 7 13'.split())for k in (1, 2, 3): print(kth_smallest(root, k))1 3 4
What this code does
- The sorted keys are 1, 3, 4, 6, 7, 8, 10, 13, 14, so these are 1, 3 and 4.
- Each call stops as soon as it has produced its key.
Example 2 · How many nodes the walk visits
class Node: def __init__(self, key): self.key = key self.left = None self.right = None
from collections import deque
def build(tokens): if not tokens or tokens[0] == "#": return None root = Node(int(tokens[0])) queue = deque([root]) rest = iter(tokens[1:]) while queue: node = queue.popleft() for side in ("left", "right"): token = next(rest, "#") if token != "#": child = Node(int(token)) setattr(node, side, child) queue.append(child) return root
visited = []
def inorder(node): if node is not None: yield from inorder(node.left) visited.append(node.key) yield node.key yield from inorder(node.right)
from itertools import islicedef kth_smallest(root, k): return next(islice(inorder(root), k - 1, None))root = build('8 3 10 1 6 # 14 # # 4 7 13'.split())print(kth_smallest(root, 2))print(visited)3 [1, 3]
What this code does
- To reach 3, the walk goes down the left side to 1, then back to 3. The right side of the tree is never entered.
- Only 2 of the 9 nodes were visited.
Time is O(h + k): go down to the smallest key, then k steps. Space is O(h). A list of all keys would cost O(n), even when k is small.
Floor and ceiling
Floor and ceiling of x are the nearest keys on each side of x. Floor is the largest key at most x; ceiling is the smallest key at least x. Either can be missing, shown as None.
Floor of x is the largest key that is at most x. Ceiling of x is the smallest key that is at least x. Search for x, and remember the last key on the path that is a candidate.
def floor(root, x): best = None node = root while node is not None: if node.key == x: return x if node.key < x: best = node.key node = node.right else: node = node.left return bestExample 3 · Floor and ceiling
class Node: def __init__(self, key): self.key = key self.left = None self.right = None
from collections import deque
def build(tokens): if not tokens or tokens[0] == "#": return None root = Node(int(tokens[0])) queue = deque([root]) rest = iter(tokens[1:]) while queue: node = queue.popleft() for side in ("left", "right"): token = next(rest, "#") if token != "#": child = Node(int(token)) setattr(node, side, child) queue.append(child) return root
def floor(root, x): best = None node = root while node is not None: if node.key == x: return x if node.key < x: best = node.key node = node.right else: node = node.left return best
def ceiling(root, x): best = None node = root while node is not None: if node.key == x: return x if node.key > x: best = node.key node = node.left else: node = node.right return best
root = build('8 3 10 1 6 # 14 # # 4 7 13'.split())for x in (5, 8, 0, 15): print(x, floor(root, x), ceiling(root, x))5 4 6 8 8 8 0 None 1 15 14 None
What this code does
- 5 is between 4 and 6, so the floor is 4 and the ceiling is 6. 8 is in the tree, so both are 8.
- 0 has no floor and 15 has no ceiling, so those show None.
A key below x is a candidate floor, and a better one can only be on the right. A key above x sends the search left. Time is O(h) and space is O(1). Ceiling is the mirror image.
Range query
To list the keys between low and high, skip the subtrees that cannot hold any. If the node is not above low, nothing to its left matters. If it is not below high, nothing to its right matters.
def keys_in_range(node, low, high): if node is None: return if low < node.key: yield from keys_in_range(node.left, low, high) if low <= node.key <= high: yield node.key if node.key < high: yield from keys_in_range(node.right, low, high)Example 4 · Keys in a range
class Node: def __init__(self, key): self.key = key self.left = None self.right = None
from collections import deque
def build(tokens): if not tokens or tokens[0] == "#": return None root = Node(int(tokens[0])) queue = deque([root]) rest = iter(tokens[1:]) while queue: node = queue.popleft() for side in ("left", "right"): token = next(rest, "#") if token != "#": child = Node(int(token)) setattr(node, side, child) queue.append(child) return root
def keys_in_range(node, low, high): if node is None: return if low < node.key: yield from keys_in_range(node.left, low, high) if low <= node.key <= high: yield node.key if node.key < high: yield from keys_in_range(node.right, low, high)
root = build('8 3 10 1 6 # 14 # # 4 7 13'.split())print(list(keys_in_range(root, 4, 10)))print(sum(keys_in_range(root, 4, 10)))print(len(list(keys_in_range(root, 20, 30))))[4, 6, 7, 8, 10] 35 0
What this code does
- The keys from 4 to 10 are 4, 6, 7, 8, 10. A generator can feed
list,sumorlendirectly. - A range with no keys gives an empty result, and the walk stops quickly.
This yields the keys in order. The cost is O(h + m) for m keys in the range, because only the nodes along the two boundary paths and the matches are visited.
A BST or a sorted list?
A sorted list answers these questions with bisect, often with less code. A BST wins when keys are inserted and deleted between queries, because a list insert moves many items.
Example 5 · The same questions with bisect
from bisect import bisect_left, bisect_right
keys = [1, 3, 4, 6, 7, 8, 10, 13, 14]print(keys[3 - 1])i = bisect_right(keys, 5)print(keys[i - 1], keys[i])print(keys[bisect_left(keys, 4):bisect_right(keys, 10)])4 4 6 [4, 6, 7, 8, 10]
What this code does
- On a sorted list the kth smallest is an index.
bisect_right(keys, 5)is the position after the last key at most 5, so the floor is just before it. - The slice from
bisect_lefttobisect_rightis the range. The list needs no tree, but inserting into it moves items.
Next step
Check your understanding
12 quick questions