Leetcode: 2476. Closest Nodes Queries in a Binary Search Tree

Problem Statement

Patterns

Prompts

Solution

The naive solution is traversing the tree to answer each query. The problem says the input is a Binary Search Tree, but there is no mention if it is balanced or not. If it is unbalanced, the naive approach would have complexity of \(O(n^2 \times m)\). This solution is too slow.

As traversing for each query is not an option, we can traverse the tree once and create a sorted array with its values (In-order Tree Traversal). The problem becomes answering the queries using the sorted array, and we can do it using a Binary Search to find the first element \(e\) that is equal or greater than the query \(q\). If \(q=a[e]\), the answer is \((q, q)\). If \(e=0\), all elements on the list are greater than \(q\) and the answer is \((-1, a[e])\). If \(e=n\), all elements on the list are smaller than \(q\) and the answer is \((a[e-1], -1)\). Otherwise, \((a[e-1],a[e])\) is the answer for the query. Time complexity is \(O(n \times m \log n)\) where \(m\) is the number of queries and \(n\) is the number of vertices. The space complexity is \(O(n)\).

class Solution:
    def closestNodes(
        self, root: Optional[TreeNode], queries: List[int]
    ) -> List[List[int]]:
        values = []

        def dfs(node):
            if node is None:
                return
            dfs(node.left)
            values.append(node.val)
            dfs(node.right)

        dfs(root)

        N = len(values)
        ans = []
        for q in queries:
            k = bisect_left(values, q)
            if k < N and values[k] == q:
                ans.append((q, q))
            elif k == 0:
                ans.append((-1, values[k]))
            elif k == N:
                ans.append((values[k - 1], -1))
            else:
                ans.append((values[k - 1], values[k]))
        return ans

Cited by 1