LeetCode in Kotlin

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

import com_github_leetcode.TreeNode

/*
 * Example:
 * var ti = TreeNode(5)
 * var v = ti.`val`
 * Definition for a binary tree node.
 * class TreeNode(var `val`: Int) {
 *     var left: TreeNode? = null
 *     var right: TreeNode? = null
 * }
 */
class Solution {
    private var k = 0
    private var count = 0
    private var `val` = 0
    fun kthSmallest(root: TreeNode, k: Int): Int {
        this.k = k
        calculate(root)
        return `val`
    }

    private fun calculate(node: TreeNode) {
        if (node.left == null && node.right == null) {
            count++
            if (count == k) {
                `val` = node.`val`
            }
            return
        }
        if (node.left != null) {
            calculate(node.left!!)
        }
        count++
        if (count == k) {
            `val` = node.`val`
            return
        }
        if (node.right != null) {
            calculate(node.right!!)
        }
    }
}