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

type TreeNode struct {
	Val   int
	Left  *TreeNode
	Right *TreeNode
}

type Node struct {
	count int
	value int
}

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func (n *Node) findkth(node *TreeNode, k int) {
	if node == nil {
		return
	}
	n.findkth(node.Left, k)
	n.count += 1
	if n.count == k {
		n.value = node.Val
		return
	}
	n.findkth(node.Right, k)
}

func kthSmallest(root *TreeNode, k int) int {
	n := &Node{
		count: 0,
		value: -1,
	}
	n.findkth(root, k)
	return n.value
}