5 Best Ways to Find the kth Smallest Element in a Binary Search Tree in Python

Rate this post

πŸ’‘ Problem Formulation: Discovering the kth smallest element in a Binary Search Tree (BST) is a common challenge that necessitates knowledge of tree traversal and data structure properties. Given a BST and an integer k, the goal is to return the kth smallest element in the tree. For instance, if the BST elements are [3, 1, 4, 2] and k=2, the expected output is 2.

Method 1: Inorder Traversal and Count

This method involves performing an inorder traversal of the BST, which visits nodes in ascending order. By maintaining a count while traversing, we can stop when we reach the kth node. This approach takes advantage of the BST property where inorder traversal gives a sorted sequence of elements.

Here’s an example:

def kth_smallest(root, k):
    stack = []
    while True:
        while root:
            stack.append(root)
            root = root.left
        root = stack.pop()
        k -= 1
        if not k:
            return root.val
        root = root.right

# Usage example
# Assume a BST is already created and its root is passed as 'root'

Output: The value of the kth smallest element in the BST.

This code snippet defines a function kth_smallest which uses an iterative approach to inorder traverse the BST using a stack, decrementing k each time a node is visited in sequence until it finds the kth smallest element.

Method 2: Recursive Inorder Traversal

In contrast with the iterative approach, recursive inorder traversal uses the call stack to keep track of visited nodes. This method is clean and requires less code but may be less efficient for large trees due to recursion depth.

Here’s an example:

def kth_smallest(root, k):
    def inorder(r):
        return inorder(r.left) + [r.val] + inorder(r.right) if r else []
    return inorder(root)[k-1]

# Usage example
# Assume a BST is already created and its root is passed as 'root'

Output: The value of the kth smallest element in the BST.

The kth_smallest function uses a nested inorder function to perform a recursive inorder traversal. It concatenates the values into a list, from which the kth smallest value is directly accessed.

Method 3: Morris Traversal Algorithm

Morris Traversal is a method that traverses the tree without using any additional space for the call stack or a stack data structure. The algorithm modifies the tree by creating temporary links known as “threads”.

Here’s an example:

def kth_smallest(root, k):
    current = root
    count = 0
    while current:
        if current.left is None:
            count += 1
            if count == k:
                return current.val
            current = current.right
        else:
            pre = current.left
            while pre.right and pre.right != current:
                pre = pre.right
            if not pre.right:
                pre.right = current
                current = current.left
            else:
                pre.right = None
                count += 1
                if count == k:
                    return current.val
                current = current.right

# Usage example
# Assume a BST is already created and its root is passed as 'root'

Output: The value of the kth smallest element in the BST.

The snippet defines a kth_smallest function implementing the Morris Traversal algorithm to find the kth smallest element. It counts elements as it traverses, leveraging the structure of the BST without additional memory overhead.

Method 4: Divide and Conquer with Node Count

The Divide and Conquer method keeps count of the number of nodes in the left subtree. This count can help us decide if the kth smallest element is in the left subtree, the right subtree, or is the current node itself.

Here’s an example:

def kth_smallest(root, k):
    def count_nodes(node):
        if node is None:
            return 0
        return 1 + count_nodes(node.left) + count_nodes(node.right)
        
    left_count = count_nodes(root.left)
    if k  left_count + 1:
        return kth_smallest(root.right, k - left_count - 1)
    return root.val

# Usage example
# Assume a BST is already created and its root is passed as 'root'

Output: The value of the kth smallest element in the BST.

This function uses a helper function count_nodes to count the nodes on the left, thus helping to find the desired kth smallest element by comparing k with the node count and choosing the right subtree to recur on.

Bonus One-Liner Method 5: Using Python’s Generator Expressions

A high-level and compact way to solve this problem is by using Python’s generator expressions in combination with inorder traversal.

Here’s an example:

def kth_smallest(root, k):
    def gen(r):
        if r is not None:
            yield from gen(r.left)
            yield r.val
            yield from gen(r.right)
    return next(x for i, x in enumerate(gen(root)) if i == k-1)

# Usage example
# Assume a BST is already created and its root is passed as 'root'

Output: The value of the kth smallest element in the BST.

This one-liner kth_smallest function leverages a generator that yields the BST’s elements in sorted order. We then use enumerate to index these elements and the next function to retrieve the kth element.

Summary/Discussion

  • Method 1: Inorder Traversal and Count. Strengths: It’s straightforward and intuitive. Weaknesses: Iterative stack could lead to higher memory usage on large trees.
  • Method 2: Recursive Inorder Traversal. Strengths: Concise and utilizes the call stack. Weaknesses: Can hit recursion limits on very large trees.
  • Method 3: Morris Traversal Algorithm. Strengths: Space efficient. Weaknesses: Modifies the tree temporarily and can be hard to understand.
  • Method 4: Divide and Conquer with Node Count. Strengths: Decides quickly which subtree to search. Weaknesses: Requires multiple traversals to count nodes, which can be inefficient.
  • Method 5: Using Python’s Generator Expressions. Strengths: Elegant and compact solution. Weaknesses: Potentially less readable for those not familiar with Python’s generators.