LeetCode in Kotlin

2476. Closest Nodes Queries in a Binary Search Tree

Medium

You are given the root of a binary search tree and an array queries of size n consisting of positive integers.

Find a 2D array answer of size n where answer[i] = [mini, maxi]:

Return the array answer.

Example 1:

Input: root = [6,2,13,1,4,9,15,null,null,null,null,null,null,14], queries = [2,5,16]

Output: [[2,2],[4,6],[15,-1]]

Explanation: We answer the queries in the following way:

Example 2:

Input: root = [4,null,9], queries = [3]

Output: [[-1,4]]

Explanation: The largest number that is smaller or equal to 3 in the tree does not exist, and the smallest number that is greater or equal to 3 is 4. So the answer for the query is [-1,4].

Constraints:

Solution

import com_github_leetcode.TreeNode
import java.util.TreeSet

/*
 * 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 val set = TreeSet<Int>()
    private fun traverse(root: TreeNode?) {
        if (root == null) {
            return
        }
        traverse(root.left)
        set.add(root.`val`)
        traverse(root.right)
    }

    fun closestNodes(root: TreeNode?, queries: List<Int>): List<List<Int>> {
        traverse(root)
        val ans: MutableList<List<Int>> = ArrayList()
        for (q in queries) {
            val list: MutableList<Int> = ArrayList()
            val min = set.floor(q)
            if (min != null) {
                list.add(min)
            } else {
                list.add(-1)
            }
            val max = set.ceiling(q)
            if (max != null) {
                list.add(max)
            } else {
                list.add(-1)
            }
            ans.add(list)
        }
        return ans
    }
}