LeetCode 230. Kth Smallest Element in a BST

Post by ailswan Nov. 08, 2023

230. Kth Smallest Element in a BST

Problem Statement

link: https://leetcode.com/problems/kth-smallest-element-in-a-bst/ https://leetcode.cn/problems/kth-smallest-element-in-a-bst/

Given the root of a binary search tree, and an integer k, return the kth smallest value (1-indexed) of all the values of the nodes in the tree.

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?

Example:

Input: root = [3,1,4,null,2], k = 1 Output: 1

Input: root = [5,3,6,2,4,null,null,1], k = 3 Output: 3

Solution Approach

The problem can be solved by performing an in-order traversal of the binary search tree (BST) and keeping track of the kth smallest element encountered.

Algorithm

Approach 1 - Recursive In-Order Traversal

  1. Recursively traverse the left subtree of the current node.
  2. Check and decrement the k value.
  3. If k becomes 0, return the current node’s value as the kth smallest element. Otherwise, continue to the right subtree. Approach 2 - Iterative In-Order Traversal using a Stack
  4. Use a stack to iteratively traverse the left subtree of the current node.
  5. Pop nodes from the stack, decrement k, and check if k becomes 0.
  6. If k is 0, return the value of the popped node as the kth smallest element. Otherwise, move to the right subtree and continue the traversal.

Implement

    class Solution:
    def kthSmallest(self, root: Optional[TreeNode], k: int) -> int:
        def inorder(root):
            if not root:
                return
            res = inorder(root.left)
            if res is not None:
                return res
            self.k -= 1
            if self.k == 0:
                return root.val
            else:
                return inorder(root.right)
            res = inorder(root.right)
            return res
        self.k = k
        return inorder(root)

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