LeetCode-in-Java

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.ArrayList;
import java.util.List;
import java.util.TreeSet;

/*
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode() {}
 *     TreeNode(int val) { this.val = val; }
 *     TreeNode(int val, TreeNode left, TreeNode right) {
 *         this.val = val;
 *         this.left = left;
 *         this.right = right;
 *     }
 * }
 */
public class Solution {
    private final TreeSet<Integer> set = new TreeSet<>();

    private void traverse(TreeNode root) {
        if (root == null) {
            return;
        }
        traverse(root.left);
        set.add(root.val);
        traverse(root.right);
    }

    public List<List<Integer>> closestNodes(TreeNode root, List<Integer> queries) {
        traverse(root);
        List<List<Integer>> ans = new ArrayList<>();
        for (int q : queries) {
            List<Integer> list = new ArrayList<>();
            Integer min = set.floor(q);
            if (min != null) {
                list.add(min);
            } else {
                list.add(-1);
            }
            Integer max = set.ceiling(q);
            if (max != null) {
                list.add(max);
            } else {
                list.add(-1);
            }
            ans.add(list);
        }
        return ans;
    }
}