LeetCode-in-All

230. Kth Smallest Element in a BST

Medium

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.

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:

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?

Solution

# Definition for a binary tree node.
# class TreeNode
#     attr_accessor :val, :left, :right
#     def initialize(val = 0, left = nil, right = nil)
#         @val = val
#         @left = left
#         @right = right
#     end
# end
# @param {TreeNode} root
# @param {Integer} k
# @return {Integer}
def kth_smallest(root, k)
  @k = k
  @count = 0
  @val = 0
  calculate(root)
  @val
end

private

def calculate(node)
  return unless node

  if node.left.nil? && node.right.nil?
    @count += 1
    @val = node.val if @count == @k
    return
  end

  calculate(node.left) if node.left
  @count += 1
  @val = node.val if @count == @k
  return if @count == @k

  calculate(node.right) if node.right
end