DSA Kth Smallest in a BST

The problem

The kth smallest value in a binary search tree is the value at position k in the tree's sorted order. Positions start at 1. The smallest value is therefore the answer for k equal to 1.

A binary search tree, or BST, keeps smaller values in each node's left subtree and larger values in its right subtree. The tree below is built by inserting nine values in the order the source lesson used.

class TreeNode:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

def insert(root, value):
    if root is None:
        return TreeNode(value)
    if value < root.value:
        root.left = insert(root.left, value)
    else:
        root.right = insert(root.right, value)
    return root

tree = None
for value in [8, 3, 10, 1, 6, 14, 4, 7, 13]:
    tree = insert(tree, value)

for k, expected in [(1, 1), (4, 6), (5, 7), (9, 14)]:
    print("k =", k, "- kth smallest:", expected)

The four questions cover the range. One asks for the smallest value. One asks for the largest. Two ask for values in the middle.

Intuition: inorder reads a BST in sorted order

An inorder walk visits a node's left subtree, then the node, then its right subtree. On a BST that order is sorted order. The BST rule guarantees it at every node.

The kth smallest is then just one position in that walk. It sits at index k minus 1 in a collected list. The walk below collects the full list and reads the position off it.

class TreeNode:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

def insert(root, value):
    if root is None:
        return TreeNode(value)
    if value < root.value:
        root.left = insert(root.left, value)
    else:
        root.right = insert(root.right, value)
    return root

tree = None
for value in [8, 3, 10, 1, 6, 14, 4, 7, 13]:
    tree = insert(tree, value)

order = []
def inorder(node):
    if node is None:
        return
    inorder(node.left)
    order.append(node.value)
    inorder(node.right)
inorder(tree)
print("inorder:", order)
print("inorder is sorted:", order == sorted(order))
print("4th smallest sits at index 3:", order[3])

The sorted check in the output is the evidence. Nothing in the walk sorts anything. The tree's own rule puts the values in this order.

Brute force: collect everything, then pick one

The brute force collects the whole inorder list first. It then returns the element at index k minus 1. A counter records the node visits the collection costs.

The method always works. Its waste is plain in the counts. Even k equal to 1, the cheapest question, pays a visit to every node in the tree.

class TreeNode:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

def insert(root, value):
    if root is None:
        return TreeNode(value)
    if value < root.value:
        root.left = insert(root.left, value)
    else:
        root.right = insert(root.right, value)
    return root

tree = None
for value in [8, 3, 10, 1, 6, 14, 4, 7, 13]:
    tree = insert(tree, value)

def kth_brute(root, k):
    order = []
    visits = [0]
    def inorder(node):
        if node is None:
            return
        inorder(node.left)
        visits[0] += 1
        order.append(node.value)
        inorder(node.right)
    inorder(root)
    return order[k - 1], visits[0]

for k in [1, 4, 9]:
    answer, visits = kth_brute(tree, k)
    print("k =", k, "- value:", answer, " nodes visited:", visits)

Nine nodes are visited for every question. The first answer is known after the first visit, yet the walk continues to the end. Stopping early is the whole improvement left on the table.

The optimal solution: count visits and stop at k

A counting inorder walk returns the value at the moment its visit count reaches k, and every caller passes that value straight back up. That one sentence is the whole method.

No list is collected. The only extra state is one counter. A found value short-circuits the rest of the walk, so nodes past the answer are never visited.

class TreeNode:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

def kth_smallest(root, k):
    visits = [0]
    def walk(node):
        if node is None:
            return None
        found = walk(node.left)
        if found is not None:
            return found
        visits[0] += 1
        if visits[0] == k:
            return node.value
        return walk(node.right)
    return walk(root), visits[0]

def insert(root, value):
    if root is None:
        return TreeNode(value)
    if value < root.value:
        root.left = insert(root.left, value)
    else:
        root.right = insert(root.right, value)
    return root

tree = None
for value in [8, 3, 10, 1, 6, 14, 4, 7, 13]:
    tree = insert(tree, value)

for k, expected in [(1, 1), (4, 6), (5, 7), (9, 14)]:
    answer, visits = kth_smallest(tree, k)
    print("k =", k, "->", answer, "expected:", expected,
          "match:", answer == expected, "visits:", visits)

All four cases from the start of the article pass. The visit count equals k every time. The largest question, k equal to 9, is the only one that pays for the whole tree.

Walkthrough: hunting the 4th smallest

The walkthrough traces the stopping walk for k equal to 4. Each line prints one counted visit. The walk ends the moment the count reaches 4.

class TreeNode:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

def insert(root, value):
    if root is None:
        return TreeNode(value)
    if value < root.value:
        root.left = insert(root.left, value)
    else:
        root.right = insert(root.right, value)
    return root

tree = None
for value in [8, 3, 10, 1, 6, 14, 4, 7, 13]:
    tree = insert(tree, value)

count = [0]
answer = [None]
def walk(node):
    if node is None or answer[0] is not None:
        return
    walk(node.left)
    if answer[0] is not None:
        return
    count[0] += 1
    print("visit", count[0], "- value", node.value)
    if count[0] == 4:
        answer[0] = node.value
        return
    walk(node.right)
walk(tree)
print("4th smallest:", answer[0])
print("nodes visited:", count[0])

The visits read 1, 3, 4, 6 in order. The fourth visit lands on 6. That value is read straight off the trace as the answer.

Five nodes of the tree are never visited. They all sit past the answer in sorted order. The trace is the proof that the walk really stops.

Complexity

The brute force visits every node and stores every value, so its time is O(n) and its extra space is O(n) for n nodes. The stopping walk counts exactly k visits before it answers, so its counted work is O(k). Its recursion stack holds one root-to-leaf path, so its extra space is O(h) for height h.

The stopping walk's advantage is a constant-factor one in the worst case, since k can equal n. Its advantage is exact for small k. The counts below show that gap at three tree sizes.

class TreeNode:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

def complete_tree(values):
    if not values:
        return None
    mid = len(values) // 2
    return TreeNode(values[mid], complete_tree(values[:mid]),
                    complete_tree(values[mid + 1:]))

def brute_visits(root):
    visits = [0]
    def inorder(node):
        if node is None:
            return
        inorder(node.left)
        visits[0] += 1
        inorder(node.right)
    inorder(root)
    return visits[0]

def stopping_visits(root, k):
    visits = [0]
    def walk(node):
        if node is None:
            return None
        found = walk(node.left)
        if found is not None:
            return found
        visits[0] += 1
        if visits[0] == k:
            return node.value
        return walk(node.right)
    walk(root)
    return visits[0]

for n in [7, 15, 31]:
    tree = complete_tree(list(range(1, n + 1)))
    print("nodes", n, "- brute visits for k = 1:", brute_visits(tree),
          " stopping visits for k = 1:", stopping_visits(tree, 1),
          " stopping visits for k = 3:", stopping_visits(tree, 3))

For k equal to 1 the brute force pays the full tree at every size. The stopping walk pays one counted visit at every size. Doubling the tree changes nothing for a small fixed k.

Edge cases and pitfalls

The first and last positions are ordinary cases. K equal to 1 answers the smallest value after one visit. K equal to 9 answers the largest value after nine visits.

Positions outside the range answer None. K equal to 0 is invalid because positions start at 1. K past the node count walks the whole tree and still finds nothing.

class TreeNode:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

def kth_smallest(root, k):
    visits = [0]
    def walk(node):
        if node is None:
            return None
        found = walk(node.left)
        if found is not None:
            return found
        visits[0] += 1
        if visits[0] == k:
            return node.value
        return walk(node.right)
    return walk(root), visits[0]

def insert(root, value):
    if root is None:
        return TreeNode(value)
    if value < root.value:
        root.left = insert(root.left, value)
    else:
        root.right = insert(root.right, value)
    return root

tree = None
for value in [8, 3, 10, 1, 6, 14, 4, 7, 13]:
    tree = insert(tree, value)

print("k = 1:", kth_smallest(tree, 1))
print("k = 9, the largest:", kth_smallest(tree, 9))
print("k = 0:", kth_smallest(tree, 0))
print("k = 10, past the end:", kth_smallest(tree, 10))
print("empty tree, k = 1:", kth_smallest(None, 1))
order = [1, 3, 4, 6, 7, 8, 10, 13, 14]
print("wrong index for k = 1, order[k]:", order[1])
print("right index for k = 1, order[k - 1]:", order[0])

The off-by-one mistake is the classic pitfall. A collected list is zero-based, so position k lives at index k minus 1. Reading index k answers the next value up, as the last two lines show.