Closest Nodes Queries in a Binary Search Tree

medium binary search tree binary search inorder traversal

Problem

Given the root of a binary search tree and an array of queries, return for each query q a pair [min, max]: min is the largest value in the tree that is ≤ q (its predecessor), and max is the smallest value in the tree that is ≥ q (its successor). Use -1 when such a value does not exist.

Inputtree = [6,2,13,1,4,9,15], queries = [2,5,16]

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

Sorted values are [1,2,4,6,9,13,15]. For 5 the predecessor is 4 and the successor is 6; for 16 nothing is ≥ 16, so max = -1.

tree (level-order, n = null)

queries (comma-separated)

def closestNodes(root, queries):
    vals = []
    def inorder(node):
        if not node:
            return
        inorder(node.left)
        vals.append(node.val)
        inorder(node.right)
    inorder(root)
    ans = []
    for q in queries:
        lo, hi = 0, len(vals)
        while lo < hi:                 # smallest idx with vals[idx] >= q
            mid = (lo + hi) // 2
            if vals[mid] < q:
                lo = mid + 1
            else:
                hi = mid
        mx = vals[lo] if lo < len(vals) else -1
        mn = vals[lo] if lo < len(vals) and vals[lo] == q else \
             (vals[lo - 1] if lo > 0 else -1)
        ans.append([mn, mx])
    return ans

function closestNodes(root, queries) {
  const vals = [];
  (function inorder(node) {
    if (!node) return;
    inorder(node.left);
    vals.push(node.val);
    inorder(node.right);
  })(root);
  const ans = [];
  for (const q of queries) {
    let lo = 0, hi = vals.length;
    while (lo < hi) {                 // smallest idx with vals[idx] >= q
      const mid = (lo + hi) >> 1;
      if (vals[mid] < q) lo = mid + 1;
      else hi = mid;
    }
    const mx = lo < vals.length ? vals[lo] : -1;
    const mn = (lo < vals.length && vals[lo] === q)
      ? vals[lo] : (lo > 0 ? vals[lo - 1] : -1);
    ans.push([mn, mx]);
  }
  return ans;
}

List<List<Integer>> closestNodes(TreeNode root, List<Integer> queries) {
    List<Integer> vals = new ArrayList<>();
    inorder(root, vals);
    List<List<Integer>> ans = new ArrayList<>();
    for (int q : queries) {
        int lo = 0, hi = vals.size();
        while (lo < hi) {              // smallest idx with vals[idx] >= q
            int mid = (lo + hi) >>> 1;
            if (vals.get(mid) < q) lo = mid + 1;
            else hi = mid;
        }
        int mx = lo < vals.size() ? vals.get(lo) : -1;
        int mn = (lo < vals.size() && vals.get(lo) == q) ? vals.get(lo)
                 : (lo > 0 ? vals.get(lo - 1) : -1);
        ans.add(List.of(mn, mx));
    }
    return ans;
}

vector<vector<int>> closestNodes(TreeNode* root, vector<int>& queries) {
    vector<int> vals;
    function<void(TreeNode*)> inorder = [&](TreeNode* n) {
        if (!n) return;
        inorder(n->left);
        vals.push_back(n->val);
        inorder(n->right);
    };
    inorder(root);
    vector<vector<int>> ans;
    for (int q : queries) {
        int lo = lower_bound(vals.begin(), vals.end(), q) - vals.begin();
        int mx = lo < (int)vals.size() ? vals[lo] : -1;
        int mn = (lo < (int)vals.size() && vals[lo] == q) ? vals[lo]
                 : (lo > 0 ? vals[lo - 1] : -1);
        ans.push_back({mn, mx});
    }
    return ans;
}

Explanation

A binary search tree has a beautiful property: its in-order traversal visits values in sorted order. So the first move is to flatten the whole tree once into a sorted array vals. After that the tree's shape is irrelevant — every query becomes a search in a sorted list.

For each query q we want two things: the predecessor (largest value ≤ q) and the successor (smallest value ≥ q). Both come from a single lower-bound binary search that finds lo, the smallest index whose value is ≥ q.

Read the answer off lo: vals[lo] is the successor (or -1 if lo ran off the end). For the predecessor, if vals[lo] equals q exactly then q itself is both answers; otherwise the predecessor sits just to the left at vals[lo - 1] (or -1 if lo is 0).

Flattening costs O(n) and each of the m queries costs O(log n), for O(n + m·log n) overall — far better than walking the tree from the root for every query when there are many queries.

Example: tree [6,2,13,1,4,9,15] flattens to [1,2,4,6,9,13,15]. Query 5 lands lo at index 3 (value 6, the successor); since 6 ≠ 5 the predecessor is vals[2] = 4, giving [4,6].

Time: O(n + m·log n) Space: O(n)