Skip to content

2476. Closest Nodes Queries in a Binary Search Tree

  • Time: $O(n + q\log n)$
  • Space: $O(n + q)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
class Solution {
 public:
  vector<vector<int>> closestNodes(TreeNode* root, vector<int>& queries) {
    vector<vector<int>> ans;
    vector<int> sortedVals;

    inorder(root, sortedVals);

    for (const int query : queries) {
      const auto it = lower_bound(sortedVals.begin(), sortedVals.end(), query);
      // query is presented in the tree, so just use {query, query}.
      if (it != sortedVals.cend() && *it == query)
        ans.push_back({query, query});
      // query isn't presented in the tree, so find the cloest one is possible.
      else
        ans.push_back({it == sortedVals.cbegin() ? -1 : *prev(it),
                       it == sortedVals.cend() ? -1 : *it});
    }

    return ans;
  }

 private:
  // Inorder walk the BST to collect sorted numbers.
  void inorder(TreeNode* root, vector<int>& sortedVals) {
    if (root == nullptr)
      return;
    inorder(root->left, sortedVals);
    sortedVals.push_back(root->val);
    inorder(root->right, sortedVals);
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
class Solution {
  public List<List<Integer>> closestNodes(TreeNode root, List<Integer> queries) {
    List<List<Integer>> ans = new ArrayList<>();
    List<Integer> sortedVals = new ArrayList<>();

    inorder(root, sortedVals);

    for (final int query : queries) {
      final int i = firstGreaterEqual(sortedVals, query);
      // query is presented in the tree, so just use {query, query}.
      if (i != sortedVals.size() && sortedVals.get(i) == query)
        ans.add(Arrays.asList(query, query));
      // query isn't presented in the tree, so find the cloest one is possible.
      else
        ans.add(Arrays.asList(i == 0 ? -1 : sortedVals.get(i - 1),
                              i == sortedVals.size() ? -1 : sortedVals.get(i)));
    }

    return ans;
  }

  // Inorder walk the BST to collect sorted numbers.
  private void inorder(TreeNode root, List<Integer> sortedVals) {
    if (root == null)
      return;
    inorder(root.left, sortedVals);
    sortedVals.add(root.val);
    inorder(root.right, sortedVals);
  }

  private int firstGreaterEqual(List<Integer> A, int target) {
    int l = 0;
    int r = A.size();
    while (l < r) {
      final int m = (l + r) / 2;
      if (A.get(m) >= target)
        r = m;
      else
        l = m + 1;
    }
    return l;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class Solution:
  def closestNodes(self, root: Optional[TreeNode], queries: List[int]) -> List[List[int]]:
    sortedVals = []
    self._inorder(root, sortedVals)

    def getClosestPair(query: int) -> List[int]:
      i = bisect_left(sortedVals, query)
      # query is presented in the tree, so just use [query, query].
      if i != len(sortedVals) and sortedVals[i] == query:
        return [query, query]
      # query isn't presented in the tree, so find the cloest one is possible.
      return [-1 if i == 0 else sortedVals[i - 1],
              -1 if i == len(sortedVals) else sortedVals[i]]

    return [getClosestPair(query) for query in queries]

  # Inorder walk the BST to collect sorted numbers.
  def _inorder(self, root: Optional[TreeNode], sortedVals: List[int]) -> None:
    if root == None:
      return
    self._inorder(root.left, sortedVals)
    sortedVals.append(root.val)
    self._inorder(root.right, sortedVals)