Kth Smallest Element in a BST
Problem
https://leetcode.com/problems/kth-smallest-element-in-a-bst/
Given the root of a binary search tree, and an integer k, return
the k:sup:`th` smallest value (1-indexed) of all
the values of the nodes in the tree.
Example 1:
Input: root = [3,1,4,null,2], k = 1
Output: 1
Example 2:
Input: root = [5,3,6,2,4,null,null,1], k = 3
Output: 3
Constraints:
The number of nodes in the tree is
n.1 <= k <= n <= 10:sup:`4`0 <= Node.val <= 10:sup:`4`
Follow up: If the BST is modified often (i.e., we can do insert and delete operations) and you need to find the kth smallest frequently, how would you optimize?
Pattern
Tree, Depth-First Search, Binary Search Tree, Binary Tree
Approaches
Explanation
Because the tree is a binary search tree, the smallest node is the leftmost
node L. To get the second smallest node, we go the right child of L
then find the leftmost node from there. If L.right doesn’t exist, then
the parent of L is the second smallest node. In general, we
go left to find the smallest node
go right, then go to 1 to find the next smallest node
if we can’t go right, go up to the parent.
We repeat the process until we reach the :math:`k`th smallest node.
To implement this, we can use a stack to keep track of the nodes visited
so far. We first push all the left children of the root onto the stack
using a while loop. We pop curr from the stack and subtract 1 from
k. If k is 0, we have found the \(k`th smallest node and can
return it. Otherwise, we move to ``curr.right`\). If curr.right exists,
we can use the previous while loop to push all the left children of
curr.right onto the stack. If curr.right = None, then the same
while loop will not push any new nodes onto the stack, and we will pop
the parent of curr on the next iteration.
Code
from __future__ import annotations
from collections import deque
class TreeNode:
"""Node in a binary tree."""
def __init__(self, val=0, left=None, right=None):
self.val = val
self.left = left
self.right = right
@classmethod
def from_list(cls, vals: list[int | None]) -> TreeNode | None:
if not vals:
return None
root = cls(vals[0])
queue = deque([root])
children = iter(vals[1:])
while queue:
node = queue.popleft()
for side in ("left", "right"):
val = next(children, None)
if val is not None:
child = cls(val)
setattr(node, side, child)
queue.append(child)
return root
def to_list(self) -> list:
result = []
queue = [self]
while queue:
node = queue.pop(0)
if node:
result.append(node.val)
queue.append(node.left)
queue.append(node.right)
else:
result.append(None)
while result and result[-1] is None:
result.pop()
return result
def kthSmallest(root: TreeNode | None, k: int) -> int:
"""Return the ``k``-th smallest value in the BST rooted at ``root``."""
stack = deque()
curr = root
while stack or curr:
while curr:
stack.append(curr)
curr = curr.left
curr = stack.pop()
k -= 1
if k == 0:
return curr.val
curr = curr.right
Test
>>> from kth_smallest_element_in_a_bst__stack import kthSmallest, TreeNode
>>> kthSmallest(TreeNode.from_list([3, 1, 4, None, 2]), 1)
1
>>> kthSmallest(TreeNode.from_list([5, 3, 6, 2, 4, None, None, 1]), 3)
3
Complexity
\(h\) is the height of the binary search tree
Measure |
Complexity |
Notes |
|---|---|---|
Time |
\(O(h + k)\) |
may need to traverse down to the tree to find the leftmost node, then visit \(k\) more nodes |
Auxiliary Space |
\(O(h)\) |
the stack may contain all nodes from the root to a leaf node |