Start learning free

LESSON 342 OF 438 · BINARY TREES

Kth Smallest Element in a BST in Python: Iterative In-Order Walk

Practice O(n) Read the full explanation

Given the root of a binary search tree (BST) and an integer k, return the k-th smallest value among all the values in the tree. k counts from 1: k = 1 asks for the smallest value.

In a BST, every value in a node's left subtree is smaller than the node's value, and every value in its right subtree is larger.

The trees below are written in level order (top to bottom, left to right), with None for a missing child.

 

Example 1:

    3
   / \
  1   4
   \
    2

Input: root = [3,1,4,None,2], k = 1
Output: 1
Explanation: In increasing order the values are 1, 2, 3, 4. The 1st smallest is 1.

Example 2:

        5
       / \
      3   6
     / \
    2   4
   /
  1

Input: root = [5,3,6,2,4,None,None,1], k = 3
Output: 3
Explanation: In increasing order the values are 1, 2, 3, 4, 5, 6. The 3rd smallest is 3.

Example 3:

Input: root = [5,3,6,2,4,None,None,1], k = 6
Output: 6
Explanation: The tree has 6 nodes, so the 6th smallest value is the largest one.

 

Constraints:

  • The number of nodes in the tree, n, is in the range [1, 1000].
  • Each node value is between -104 and 104, and all values are different.
  • 1 <= k <= n

 

Follow up: Can you stop as soon as you reach the answer, in O(h + k) time, where h is the height of the tree?
STEP 1 OF 12
    Ctrl + Enter to check

    PYTHON · BINARY TREES

    Kth Smallest Element In Bst in Python, explained

    Problem overview

    The method kthSmallest receives the root of a binary search tree and an integer k, counted from 1. It returns the k-th smallest value in the tree, so k = 1 asks for the smallest. All values are different, and k is never larger than the number of nodes.

    For the tree with root 3, children 1 and 4, and 2 as the right child of 1, the values in increasing order are 1, 2, 3, 4. With k = 1 the answer is 1. For the tree with root 5, children 3 and 6, then 2 and 4 under 3 and 1 under 2, the sorted values are 1, 2, 3, 4, 5, 6. With k = 3 the answer is 3, and with k = 6 it is 6.

    The follow-up asks you to stop as soon as the answer is found, in O(h + k) time where h is the height of the tree.

    Approach: in-order traversal, stopped early

    In a BST, everything in a node's left subtree is smaller than the node, and everything in its right subtree is larger. An in-order traversal visits the left subtree, then the node, then the right subtree, so it visits the values in increasing order. The k-th node it visits holds the answer.

    The simple way is to collect all values in a list and take element k - 1, which costs O(n) time and memory even when k is 1. A recursive traversal can't easily stop in the middle. So this lesson does the traversal with an explicit stack, which lets the function return the moment it reaches the k-th node and leaves the rest of the tree unvisited.

    The stack holds the nodes on the way down that have been passed but not yet visited. To find the next smallest value, go as far left as possible, pushing every node on the way. The top of the stack is then the smallest value not yet visited. Visit it, then continue with its right subtree, which holds the values just above it and below the nodes still waiting on the stack.

    On the second example with k = 3: going left from 5 pushes 5, 3, 2 and 1. Popping gives 1 (count 1). Its right child is empty, so the next pop gives 2 (count 2), and then 3 (count 3), which is returned.

    Step by step

    1. Create an empty list stack. A Python list works as a stack with append and pop at its end.
    2. Create node at root. It is the root of the part of the tree still to be walked.
    3. Create count at 0. It counts the nodes visited so far.
    4. Loop while node or stack. The walk is over only when there is no subtree left to enter and nothing waiting on the stack.
    5. Go left with an inner loop that runs while node is not None.
    6. Push node inside that loop. It has to wait until the smaller values in its left subtree are visited.
    7. Move node to node.left.
    8. Pop the top of the stack into node. Everything smaller than it has already been visited.
    9. Add 1 to count. The node is visited now.
    10. Check count == k. If so, return node.val and stop.
    11. Move to node.right. The next larger values are in the right subtree. If it is empty, node becomes None and the next pass pops the next waiting node.
    12. Return -1 after the loop. Because k is at most the number of nodes, the function returns earlier, so this line is never reached.

    The solution in Python

    class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        stack = []
        node = root
        count = 0
        while node or stack:
            while node:
                stack.append(node)
                node = node.left
            node = stack.pop()
            count += 1
            if count == k:
                return node.val
            node = node.right
        return -1

    The Python details worth a closer look:

    • Truthiness. A node object is truthy and None is falsy, and an empty list is falsy too. That is why while node or stack and while node work without explicit comparisons.
    • pop() without an argument removes and returns the last element, in constant time. It raises IndexError on an empty list, but the outer condition guarantees that the stack isn't empty when node is None.
    • The variable node has two jobs. It first walks down the left spine, ending at None, and then it is reassigned to the popped node. That reuse is deliberate: after each pop, node.right decides where to go next.
    • No recursion limit. Python limits recursion depth, by default to about 1000 calls. This version keeps its own stack in a list, so a long chain of nodes can't raise RecursionError.

    Complexity

    Time: O(h + k). The first descent pushes at most h nodes. After that, each visited node costs a pop and some pushes, and every node is pushed and popped at most once. The walk stops at the k-th visit, so it does k pops, plus the pushes needed to reach them, which are at most h + k in total. In the worst case, k = n, this is O(n).

    Space: O(h). The stack holds at most one path from the root down, so its size is bounded by the height. For a balanced tree that is about log n, and for a chain it is n.

    Edge cases and common mistakes

    • Counting from 0. k starts at 1, which is why count is increased before the comparison with k.
    • A single node. With one node and k = 1, the left spine pushes the node, the pop visits it, and its value is returned.
    • The largest value. With k equal to the number of nodes, the whole tree is visited and the last pop returns the maximum.
    • Forgetting to go right. Without node = node.right, only the left spine and the nodes on the stack would be visited, and right subtrees would be skipped.
    • Checking the count too early. The count must grow when a node is popped, not when it is pushed. Pushing happens in the wrong order for sorted output.
    • Collecting everything first. Building a full sorted list works, but it ignores the early exit that the follow-up asks for.
    • A left-leaning chain. For [3, 2, None, 1] and k = 1, the first descent pushes all three nodes before the first visit, so the stack depth is the full height.

    Testing

    • root = [3, 1, 4, None, 2], k = 1 gives 1.
    • root = [5, 3, 6, 2, 4, None, None, 1], k = 3 gives 3.
    • The same tree with k = 6 gives 6.
    • The same tree with k = 1 gives 1, and with k = 5 gives 5.
    Keep going with PythonAbout 340 step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free