230. Kth Smallest Element in a BST

• categories: Code, Interview Question, leetcode, Medium
• source: https://leetcode.com/problems/kth-smallest-element-in-a-bst
• topics: Tree

Description

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.

Idea

Inorder traversal with counter and exit if answer is found

Code

Inorder traversal with recursion

class Solution:
    def kthSmallest(self, root: Optional[TreeNode], k: int) -> int:
        i = 0
        ans = None
        def inorder(node):
            nonlocal i, ans
            if ans: 
                return
            if node.left:
                inorder(node.left)
            i += 1
            if i == k:
                ans = node.val
                return
            
            if node.right:
                inorder(node.right)
 
        inorder(root)
        return ans 

With generator, more pytonic

class Solution:
    def kthSmallest(self, root: Optional[TreeNode], k: int) -> int:
        def in_order_traversal(node):
            if node:
                yield from in_order_traversal(node.left)
                yield node.val
                yield from in_order_traversal(node.right)
        
        gen = in_order_traversal(root)
        for _ in range(k):
            result = next(gen)
        return result