Kth Smallest Element in a BST in Python

Finding the Kth smallest element in a Binary Search Tree (BST) is a common problem that leverages the BST property where in-order traversal visits nodes in sorted order. We can use in-order traversal to collect elements and return the Kth smallest one.

For K=3, the answer is 7 since in-order traversal gives us: 2, 5, 7, 10, 13, 15.

Algorithm Steps

The solution uses in-order traversal to visit nodes in ascending order ?

• Create an empty list to store node values
• Perform in-order traversal: left subtree ? root ? right subtree
• Return the (K-1)th element from the sorted list

Implementation

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

def insert(temp, data):
    queue = [temp]
    while queue:
        node = queue.pop(0)
        if not node.left:
            node.left = TreeNode(data)
            break
        else:
            queue.append(node.left)
        if not node.right:
            node.right = TreeNode(data)
            break
        else:
            queue.append(node.right)

def make_tree(elements):
    tree = TreeNode(elements[0])
    for element in elements[1:]:
        insert(tree, element)
    return tree

class Solution:
    def kthSmallest(self, root, k):
        nodes = []
        self.inorder_traversal(root, nodes)
        return nodes[k-1]
    
    def inorder_traversal(self, root, nodes):
        if root is None:
            return
        # Traverse left subtree
        self.inorder_traversal(root.left, nodes)
        # Visit root
        nodes.append(root.data)
        # Traverse right subtree
        self.inorder_traversal(root.right, nodes)

# Example usage
solution = Solution()
tree = make_tree([10, 5, 15, 2, 7, 13])
result = solution.kthSmallest(tree, 3)
print(f"3rd smallest element: {result}")

3rd smallest element: 7

Optimized Approach

Instead of storing all elements, we can stop traversal once we find the Kth element ?

class OptimizedSolution:
    def kthSmallest(self, root, k):
        self.count = 0
        self.result = None
        self.inorder(root, k)
        return self.result
    
    def inorder(self, root, k):
        if root is None or self.result is not None:
            return
        
        # Traverse left subtree
        self.inorder(root.left, k)
        
        # Process current node
        self.count += 1
        if self.count == k:
            self.result = root.data
            return
        
        # Traverse right subtree
        self.inorder(root.right, k)

# Test the optimized solution
optimized = OptimizedSolution()
tree = make_tree([10, 5, 15, 2, 7, 13])
result = optimized.kthSmallest(tree, 3)
print(f"3rd smallest element (optimized): {result}")

3rd smallest element (optimized): 7

Comparison

Where n is the number of nodes and h is the height of the tree.

Conclusion

The in-order traversal approach leverages BST properties to find the Kth smallest element efficiently. Use the optimized version when dealing with large trees and small K values for better performance.