Ordered maps: floor, ceiling and rank queries

"""Floor and ceiling in a binary search tree: one walk from the root, remembering the best candidate so far."""


class Node:
    __slots__ = ("key", "left", "right")

    def __init__(self, key):
        self.key, self.left, self.right = key, None, None


def insert(root, key):
    if root is None:
        return Node(key)
    if key < root.key:
        root.left = insert(root.left, key)
    elif key > root.key:
        root.right = insert(root.right, key)
    return root


def floor(root, x, path):
    """The largest key <= x: a key <= x is a candidate, and anything better can only be to its right."""
    best, node = None, root
    while node:
        path.append(node.key)
        if node.key == x:
            return x
        if node.key < x:
            best, node = node.key, node.right
        else:
            node = node.left
    return best


def ceiling(root, x, path):
    """The smallest key >= x: the mirror image."""
    best, node = None, root
    while node:
        path.append(node.key)
        if node.key == x:
            return x
        if node.key > x:
            best, node = node.key, node.left
        else:
            node = node.right
    return best


root = None
for key in (120, 60, 180, 30, 75, 135, 240, 15):  # the free slots, inserted in a balanced order
    root = insert(root, key)
for x in (40, 130, 250, 5):
    fp, cp = [], []
    f, c = floor(root, x, fp), ceiling(root, x, cp)
    print(f"x = {x:3}: floor {f}, path {fp}; ceiling {c}, path {cp}")

x =  40: floor 30, path [120, 60, 30]; ceiling 60, path [120, 60, 30]
x = 130: floor 120, path [120, 180, 135]; ceiling 135, path [120, 180, 135]
x = 250: floor 240, path [120, 180, 240]; ceiling None, path [120, 180, 240]
x =   5: floor None, path [120, 60, 30, 15]; ceiling 15, path [120, 60, 30, 15]

"""Free appointment slots in a sorted list: floor, ceiling and a range, with bisect."""

from bisect import bisect_left, bisect_right, insort

free = [15, 30, 60, 75, 120, 135, 180, 240]  # minutes after 09:00 that are still free


def clock(minutes):
    return f"{9 + minutes // 60:02}:{minutes % 60:02}"


def ceiling(a, x):
    """The smallest value >= x, or None."""
    i = bisect_left(a, x)  # the first position whose value is >= x
    return a[i] if i < len(a) else None


def floor(a, x):
    """The largest value <= x, or None."""
    i = bisect_right(a, x)  # the first position whose value is > x
    return a[i - 1] if i > 0 else None


def show(minutes):
    return clock(minutes) if minutes is not None else "none"


for wanted in (40, 60, 250, 5):
    print(f"wants {clock(wanted)}: at or after {show(ceiling(free, wanted))}, at or before {show(floor(free, wanted))}")

lo, hi = bisect_left(free, 60), bisect_right(free, 180)  # every free slot from 10:00 to 12:00, both included
print("free from 10:00 to 12:00:", ", ".join(clock(m) for m in free[lo:hi]))

del free[bisect_left(free, 60)]  # 10:00 is booked
insort(free, 45)  # a cancellation frees 09:45
print("after booking 10:00 and freeing 09:45, a patient who wants 10:00 gets", show(ceiling(free, 60)),
      "or", show(floor(free, 60)))

wants 09:40: at or after 10:00, at or before 09:30
wants 10:00: at or after 10:00, at or before 10:00
wants 13:10: at or after none, at or before 13:00
wants 09:05: at or after 09:15, at or before none
free from 10:00 to 12:00: 10:00, 10:15, 11:00, 11:15, 12:00
after booking 10:00 and freeing 09:45, a patient who wants 10:00 gets 10:15 or 09:45

"""Rank and k-th smallest from subtree sizes, checked against bisect on 100,000 shuffled scores."""

import random
from bisect import bisect_left


class Node:
    __slots__ = ("key", "left", "right", "size")

    def __init__(self, key):
        self.key, self.left, self.right, self.size = key, None, None, 1


def size(node):
    return node.size if node else 0


def insert(root, key):
    """Ordinary BST insertion that also adds 1 to the size of every node it passes (keys are distinct)."""
    if root is None:
        return Node(key)
    node = root
    while True:
        node.size += 1
        side = "left" if key < node.key else "right"
        child = getattr(node, side)
        if child is None:
            setattr(node, side, Node(key))
            return root
        node = child


def rank(root, x):
    """How many keys are smaller than x: every time the walk goes right, the left subtree and the node count."""
    smaller, node, visited = 0, root, 0
    while node:
        visited += 1
        if x <= node.key:
            node = node.left
        else:
            smaller += size(node.left) + 1
            node = node.right
    return smaller, visited


def select(root, k):
    """The k-th smallest key, counting from 0: compare k with the size of the left subtree at each node."""
    node, visited = root, 0
    while node:
        visited += 1
        left = size(node.left)
        if k < left:
            node = node.left
        elif k == left:
            return node.key, visited
        else:
            k -= left + 1
            node = node.right
    raise IndexError("k is out of range")


rng = random.Random(8)
scores = rng.sample(range(10_000_000), 100_000)  # distinct scores, in random order
root = None
for s in scores:
    root = insert(root, s)
ordered = sorted(scores)

print(f"rank of 5,000,000: {rank(root, 5_000_000)[0]:,} keys are smaller; bisect says {bisect_left(ordered, 5_000_000):,}")
print("select k = 0, 49,999 and 99,999:", ", ".join(f"{select(root, k)[0]:,}" for k in (0, 49_999, 99_999)))
print("               bisect says     ", ", ".join(f"{ordered[k]:,}" for k in (0, 49_999, 99_999)))

worst = 0
for _ in range(10_000):  # random checks against the sorted list
    x = rng.randrange(10_000_000)
    r, v1 = rank(root, x)
    k = rng.randrange(100_000)
    key, v2 = select(root, k)
    assert r == bisect_left(ordered, x) and key == ordered[k]
    worst = max(worst, v1, v2)
print(f"10,000 random rank and select queries agree with bisect; the longest walk visited {worst} nodes")

rank of 5,000,000: 49,973 keys are smaller; bisect says 49,973
select k = 0, 49,999 and 99,999: 98, 5,002,295, 9,999,793
               bisect says      98, 5,002,295, 9,999,793
10,000 random rank and select queries agree with bisect; the longest walk visited 37 nodes