Kth smallest with subtree sizes in Python

Why it matters

Some questions ask for a position in sorted order, not for a value. "Who is the third fastest runner?" "What is the median response time?" "Which score is at rank 100 on the leaderboard?" Each of these asks for the k-th value in order, where k is a position counted from 1.

The naive answer is to walk the whole tree in order and count until you reach position k. That takes O(n) steps, where n is the number of stored values. A leaderboard with a million players would need a million steps for each rank query.

This lesson stores one extra number in each node, the size of the subtree below it. With that number, a query for the k-th smallest value walks one path from the root and takes O(height) steps. The floor and ceiling lesson used the same walk from the root. Here the walk also uses the sizes to decide which way to go.

By the end of this lesson you will be able to write kth_smallest and kth_largest for a Python BST, update the sizes on insert and delete, and check the sizes with a small test.

Start with the problem

Suppose the stored values are 20, 30, 35, 40, 50, 60, 70, 80. Their sorted order is the same as the list. The 5th smallest is 50, and the 2nd largest is 70.

You can answer both by writing the list out in order and counting. That is easy with eight values. The point of a tree is to avoid the count.

A sorted Python list answers the k-th query in one step, but adding a value to it can shift many items. The sizes in this lesson make the k-th query fast on a tree, and inserts still take O(height) steps.

The idea: count the left side

Every node gets one more number, called size. The size of a node is the number of nodes in its subtree, counting the node itself. A leaf has size 1.

Here is the tree for the eight values above, with sizes in brackets. The tree was built by inserting 50, 30, 70, 20, 40, 60, 80 and 35, in that order.

              50 [8]
            /        \
       30 [4]        70 [3]
      /      \       /      \
  20 [1]   40 [2]  60 [1]   80 [1]
           /
        35 [1]

Check a few sizes by hand. Node 40 has one child, 35, so its size is 1 + 1 = 2. Node 30 has 20 on the left (size 1) and 40 on the right (size 2), so its size is 1 + 1 + 2 = 4. The root has 30 on the left (size 4) and 70 on the right (size 3), so its size is 1 + 4 + 3 = 8.

Now the key fact. In a BST, every value in the left subtree is smaller than the node, and every value in the right subtree is larger. So the size of the left subtree tells you exactly how many values come before the node in sorted order.

At the root, the left subtree has 4 values. So the root is the 5th smallest. That is exactly 50, the 5th value in the sorted list.

The walk for kth_smallest

Let remaining be the position you still need. Start at the root with remaining = k. At each node, let left_size be the size of its left subtree. There are three cases.

remaining is less than or equal to left_size. The answer is among the smaller values, so go left. remaining does not change.
remaining is exactly left_size + 1. The node is the answer. Return its key.
remaining is larger than left_size + 1. The answer is on the right. Skip the left subtree and the node itself, which together hold left_size + 1 values. Subtract that from remaining and go right.

Each step goes one level down, so the walk takes at most as many steps as the tree has levels. That is O(height), the same cost as a floor or ceiling query.

For kth_largest, you do not need a second walk. The k-th largest is the same as the (n minus k plus 1)-th smallest, where n is the number of values. The 1st largest is the 8th smallest in an eight-value tree, and the 2nd largest is the 7th smallest.

Worked example 1: the 5th smallest

Find kth_smallest(5) on the tree above.

Node	left_size	remaining	Action
50	4	5	5 equals 4 + 1, so return 50

The answer is 50 in one step. The left subtree held 4 values, and the root was the next one in order.

Worked example 2: the 6th smallest

Find kth_smallest(6).

Node	left_size	remaining	Action
50	4	6	6 is larger than 5, so subtract 5 and go right; remaining is 1
70	1	1	1 is at most 1, so go left
60	0	1	1 equals 0 + 1, so return 60

The answer is 60. Check it against the sorted list: 20, 30, 35, 40, 50, 60, and the 6th value is 60.

Notice that the walk never looked at 20, 30, 35 or 40. Those values were all skipped by the size of the left subtree at 50. That is the whole speed-up.

Worked example 3: kth largest

Find the 2nd largest value, kth_largest(2). With 8 values, this is kth_smallest(8 - 2 + 1), which is kth_smallest(7).

Node	left_size	remaining	Action
50	4	7	7 is larger than 5, so subtract 5 and go right; remaining is 2
70	1	2	2 equals 1 + 1, so return 70

The answer is 70. In the sorted list 20 to 80, the 2nd largest is 70.

Keeping the sizes right on insert

When you insert a value, every node on the path from the root to the new leaf gets one more value below it. Their sizes must grow by one.

There are two ways to do this. The first adds 1 to each node on the path. The second recomputes each size from its children: the size is 1 plus the size of the left child plus the size of the right child.

The code below uses the second way. The insert function calls itself on the left or right child, and when that call returns, the current node recomputes its own size. Because each node on the path is visited on the way back up, the sizes are fixed from the bottom to the top.

Insert 45 into the tree above. The path is 50, then 30, then 40, where 45 becomes the new right child. On the way back up, 40 recomputes to 3, 30 recomputes to 5, and 50 recomputes to 9. The root now has size 9, which is correct for nine values.

Inserting a value that is already present changes nothing, and the recomputed sizes are still correct.

Keeping the sizes right on delete

Delete has three cases, and the sizes must be right after each.

The node has no left child. Replace the node with its right child. The right child carries its own size.
The node has no right child. Replace the node with its left child.
The node has two children. Find the smallest key in the right subtree, which is the successor. Copy its key into the node, then delete the successor from the right subtree. The successor has no left child, so this second delete is case 1 or case 2.

After any of these, the nodes on the path from the root to the removed node have one fewer value below them. Recomputing each size on the way back up fixes all of them.

Delete 30 from the tree above. Node 30 has two children, 20 and 40. Its successor is the smallest key in the right subtree, which is 35, the left child of 40. Node 30 takes the key 35. Then 35 is removed from the right subtree, so node 40 loses its left child and its size becomes 1. The node that now holds 35 has size 3. The root recomputes to 7, which is correct for the seven values left.

The code

The program below is the full tested file. It has the insert, delete and query functions, a check that every size is right, and a random test against a sorted list of the stored values.

Python · editable

import random

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

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

def update(node):
    node.size = 1 + size(node.left) + size(node.right)

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)
    update(root)
    return root

def delete(root, key):
    if root is None:
        return None
    if key < root.key:
        root.left = delete(root.left, key)
    elif key > root.key:
        root.right = delete(root.right, key)
    else:
        if root.left is None:
            return root.right
        if root.right is None:
            return root.left
        successor = root.right
        while successor.left is not None:
            successor = successor.left
        root.key = successor.key
        root.right = delete(root.right, successor.key)
    update(root)
    return root

def kth_smallest(root, k):
    if k < 1 or k > size(root):
        return None
    cur = root
    remaining = k
    while cur is not None:
        left_size = size(cur.left)
        if remaining <= left_size:
            cur = cur.left
        elif remaining == left_size + 1:
            return cur.key
        else:
            remaining -= left_size + 1
            cur = cur.right
    return None

def kth_largest(root, k):
    if k < 1 or k > size(root):
        return None
    return kth_smallest(root, size(root) - k + 1)

def sizes_are_right(node):
    if node is None:
        return True
    expected = 1 + size(node.left) + size(node.right)
    return node.size == expected and sizes_are_right(node.left) and sizes_are_right(node.right)

def inorder(node, out):
    if node is None:
        return
    inorder(node.left, out)
    out.append(node.key)
    inorder(node.right, out)

def show(value):
    return "none" if value is None else str(value)

def check_random(rounds):
    rnd = random.Random(7)
    checks = 0
    mismatches = 0
    bad_sizes = 0
    for _ in range(rounds):
        root = None
        reference = set()
        for _ in range(200):
            v = rnd.randrange(100)
            if rnd.randrange(3) == 0:
                root = delete(root, v)
                reference.discard(v)
            else:
                root = insert(root, v)
                reference.add(v)
            if not sizes_are_right(root):
                bad_sizes += 1
            ordered = sorted(reference)
            for k in range(0, len(ordered) + 2):
                checks += 1
                expected = ordered[k - 1] if 1 <= k <= len(ordered) else None
                if kth_smallest(root, k) != expected:
                    mismatches += 1
    print(f"kth checks={checks} mismatches={mismatches} trees with wrong sizes={bad_sizes}")

root = None
for key in [50, 30, 70, 20, 40, 60, 80, 35]:
    root = insert(root, key)
keys = []
inorder(root, keys)
print("in order:", keys)
print("size of root:", size(root))
for k in range(1, 10, 4):
    print(f"kth_smallest({k})={show(kth_smallest(root, k))} kth_largest({k})={show(kth_largest(root, k))}")
print("kth_smallest(0)=" + show(kth_smallest(root, 0)))

root = delete(root, 30)
keys = []
inorder(root, keys)
print("after deleting 30, in order:", keys)
print("size of root:", size(root), ", sizes right:", sizes_are_right(root))
print("kth_smallest(3)=" + show(kth_smallest(root, 3)) + " kth_largest(2)=" + show(kth_largest(root, 2)))
check_random(20)

The real output is:

in order: [20, 30, 35, 40, 50, 60, 70, 80]
size of root: 8
kth_smallest(1)=20 kth_largest(1)=80
kth_smallest(5)=50 kth_largest(5)=40
kth_smallest(9)=none kth_largest(9)=none
kth_smallest(0)=none
after deleting 30, in order: [20, 35, 40, 50, 60, 70, 80]
size of root: 7 , sizes right: True
kth_smallest(3)=40 kth_largest(2)=70
kth checks=155862 mismatches=0 trees with wrong sizes=0

The first block matches the eight values in order. kth_largest(5) is 40, which is the 4th smallest, because 8 minus 5 plus 1 is 4. After the delete, the tree holds seven values and the 3rd smallest is 40.

The last line comes from check_random. It runs 20 random trees. Each tree gets 200 random inserts and deletes. After every change it checks that all sizes are right, and it compares every kth_smallest answer with the sorted list of stored values. The random numbers come from a seeded generator, so the count is the same on every run on the same Python version.

Reading the Python code

size returns 0 for None. This matters because a leaf's children are None, and reading .size from None stops the program. The mistakes section shows that error.
update is the one place where a size is computed. Every insert and delete ends by calling it on the node that changed.
In delete, the two-child case copies the successor's key into the node, then deletes the successor from the right side. The node is not replaced, so the update call at the end still runs on it.
kth_smallest returns None for k below 1 or above the tree size. The check is a shortcut. Without it, the walk would still end at an empty child and return None, but it would do extra work.

Common mistakes

Forgetting to update sizes on delete

The most common mistake is to update sizes on insert and forget delete. The insert code still looks right, so the bug can hide for a long time.

The code below is a copy of the delete function with the update(root) line removed. The rest of the code is the same. It uses the Node class and the other functions from the run block above, saved as sized_bst.py in the same folder.

from sized_bst import Node, insert, size, sizes_are_right, kth_smallest, show
 
 
def delete_no_update(root, key):
    if root is None:
        return None
    if key < root.key:
        root.left = delete_no_update(root.left, key)
    elif key > root.key:
        root.right = delete_no_update(root.right, key)
    else:
        if root.left is None:
            return root.right
        if root.right is None:
            return root.left
        successor = root.right
        while successor.left is not None:
            successor = successor.left
        root.key = successor.key
        root.right = delete_no_update(root.right, successor.key)
    return root
 
 
root = None
for key in [50, 30, 70, 20, 40, 60, 80, 35]:
    root = insert(root, key)
root = delete_no_update(root, 30)
print("sizes right after bad delete:", sizes_are_right(root))
print("root size reported:", size(root))
print("kth_smallest(7)=" + show(kth_smallest(root, 7)))
print("kth_smallest(8)=" + show(kth_smallest(root, 8)))

Importing the module also runs its run block, so the output starts with the ten lines that block prints. The real output ends with these four lines:

sizes right after bad delete: False
root size reported: 8
kth_smallest(7)=70
kth_smallest(8)=80

After the bad delete the tree holds seven values, but the root still says 8. The 7th smallest is reported as 70, and the correct answer is 80. The 8th smallest is reported as 80, but the tree has only seven values, so the answer should be None. A wrong size gives a wrong answer with no error message, which is why the sizes_are_right check matters in tests.

Reading the size of a `None` node

A helper that reads node.size without checking for None stops the program when it gets None. Here is a one-line helper that does that, called on purpose:

def bad_size(node):
    return node.size
 
print(bad_size(None))

The last line of the output is:

AttributeError: 'NoneType' object has no attribute 'size'

The message names the attribute and the type of the value. The fix is the same as in size: return 0 for None.

Off by one in the walk

The two comparisons in the walk are easy to get wrong. The left case must use remaining <= left_size, not remaining < left_size. With <, the value at position left_size would be skipped, because it belongs to the left side.

The right case must subtract left_size + 1, not left_size. The node itself is one of the values skipped, so it has to be counted.

Wrong formula for kth largest

A common error is kth_smallest(root, size(root) - k), which is one short. On the eight-value tree, kth_largest(1) should be 80. With the wrong formula, kth_smallest(7) is called, which returns 70. The correct formula is size(root) - k + 1, which for k equal to 1 gives 8.

Starting the count at zero

The positions in this lesson start at 1. kth_smallest(1) is the smallest value, and kth_smallest(0) has no answer. Code that copies a zero-based list index into this function will ask for the wrong value by one. The code returns None for k equal to 0, which makes the mistake easy to see.

Practice

Exercise 1. Insert 10, 5, 15 and 12, in that order, into an empty tree. Give the size of each node.

Show the answer

The root 10 has size 4. The node 5 has size 1. The node 15 has size 2, because its left child 12 is below it. The node 12 has size 1.

Exercise 2. Use the tree from Exercise 1. Find kth_smallest(3) by walking the tree.

Show the answer

The answer is 12. At 10, left_size is 1. Three is larger than 1 + 1 = 2, so subtract 2, which leaves 1, and go right to 15. At 15, left_size is 1, and the remaining position 1 is at most 1, so go left to 12. At 12, left_size is 0, and 1 equals 0 + 1, so return 12.

Exercise 3. Use the tree from Exercise 1. Delete 10, using the two-child rule. Give the new root and the sizes of all three nodes.

Show the answer

The successor of 10 is the smallest key in the right subtree, which is 12. The root takes the key 12. Then 12 is removed from the right subtree, which leaves 15 with no left child. The new tree has root 12, with 5 on the left and 15 on the right. The root has size 3, node 5 has size 1 and node 15 has size 1. The in-order list is 5, 12, 15, which is correct.

Exercise 4. A teammate writes kth_largest as kth_smallest(root, size(root) - k). On the eight-value tree from this lesson, they test k equal to 1 and expect 80. What do they get, and why?

Show the answer

They get 70. The formula calls kth_smallest(7), which is the 7th value in order, and that is 70. The largest value is the 8th smallest, so the formula must be size(root) - k + 1. With that formula, k equal to 1 becomes kth_smallest(8), which returns 80.

Summary

Storing the size of each subtree lets a walk from the root skip whole subtrees while it looks for the k-th value.
The size of the left subtree tells how many values come before a node, so the walk compares remaining with it.
kth_largest(k) is kth_smallest(size - k + 1), so one walk serves both queries.
Recompute each node's size from its children on the way back up after every insert and delete.
A wrong size gives wrong answers without any error, so test with a function that checks every size and compares the answers with a sorted reference.

Kth smallest with subtree sizes in Python

Why it matters

Start with the problem

The idea: count the left side

The walk for kth_smallest

Worked example 1: the 5th smallest

Worked example 2: the 6th smallest

Worked example 3: kth largest

Keeping the sizes right on insert

Keeping the sizes right on delete

The code

Reading the Python code

Common mistakes

Forgetting to update sizes on delete

Reading the size of a None node

Off by one in the walk

Wrong formula for kth largest

Starting the count at zero

Practice

Summary

Stuck? Ask the tutor

Reading the size of a `None` node