3275. K-th Nearest Obstacle Queries

Description

There is an infinite 2D plane.

You are given a positive integer k. You are also given a 2D array queries, which contains the following queries:

queries[i] = [x, y]: Build an obstacle at coordinate (x, y) in the plane. It is guaranteed that there is no obstacle at this coordinate when this query is made.

After each query, you need to find the distance of the k^th nearest obstacle from the origin.

Return an integer array results where results[i] denotes the k^th nearest obstacle after query i, or results[i] == -1 if there are less than k obstacles.

Note that initially there are no obstacles anywhere.

The distance of an obstacle at coordinate (x, y) from the origin is given by |x| + |y|.

Example 1:

Input: queries = [[1,2],[3,4],[2,3],[-3,0]], k = 2

Output: [-1,7,5,3]

Explanation:

Initially, there are 0 obstacles.
After queries[0], there are less than 2 obstacles.
After queries[1], there are obstacles at distances 3 and 7.
After queries[2], there are obstacles at distances 3, 5, and 7.
After queries[3], there are obstacles at distances 3, 3, 5, and 7.

Example 2:

Input: queries = [[5,5],[4,4],[3,3]], k = 1

Output: [10,8,6]

Explanation:

After queries[0], there is an obstacle at distance 10.
After queries[1], there are obstacles at distances 8 and 10.
After queries[2], there are obstacles at distances 6, 8, and 10.

Constraints:

1 <= queries.length <= 2 * 10⁵
All queries[i] are unique.
-10⁹ <= queries[i][0], queries[i][1] <= 10⁹
1 <= k <= 10⁵

Solutions

Solution 1: Priority Queue (Max-Heap)

We can use a priority queue (max-heap) to maintain the $k$ obstacles closest to the origin.

Traverse $\textit{queries}$, and for each query, calculate the sum of the absolute values of $x$ and $y$, then add it to the priority queue. If the size of the priority queue exceeds $k$, pop the top element. If the current size of the priority queue is equal to $k$, add the top element to the answer array; otherwise, add $-1$ to the answer array.

After the traversal, return the answer array.

The time complexity is $O(n \times \log k)$, and the space complexity is $O(k)$. Here, $n$ is the length of the array $\textit{queries}$.

Python3JavaC++GoTypeScript

class Solution:
    def resultsArray(self, queries: List[List[int]], k: int) -> List[int]:
        ans = []
        pq = []
        for i, (x, y) in enumerate(queries):
            heappush(pq, -(abs(x) + abs(y)))
            if i >= k:
                heappop(pq)
            ans.append(-pq[0] if i >= k - 1 else -1)
        return ans

class Solution {
    public int[] resultsArray(int[][] queries, int k) {
        int n = queries.length;
        int[] ans = new int[n];
        PriorityQueue<Integer> pq = new PriorityQueue<>(Collections.reverseOrder());
        for (int i = 0; i < n; ++i) {
            int x = Math.abs(queries[i][0]) + Math.abs(queries[i][1]);
            pq.offer(x);
            if (i >= k) {
                pq.poll();
            }
            ans[i] = i >= k - 1 ? pq.peek() : -1;
        }
        return ans;
    }
}

class Solution {
public:
    vector<int> resultsArray(vector<vector<int>>& queries, int k) {
        vector<int> ans;
        priority_queue<int> pq;
        for (const auto& q : queries) {
            int x = abs(q[0]) + abs(q[1]);
            pq.push(x);
            if (pq.size() > k) {
                pq.pop();
            }
            ans.push_back(pq.size() == k ? pq.top() : -1);
        }
        return ans;
    }
};

func resultsArray(queries [][]int, k int) (ans []int) {
    pq := &hp{}
    for _, q := range queries {
        x := abs(q[0]) + abs(q[1])
        pq.push(x)
        if pq.Len() > k {
            pq.pop()
        }
        if pq.Len() == k {
            ans = append(ans, pq.IntSlice[0])
        } else {
            ans = append(ans, -1)
        }
    }
    return
}

func abs(x int) int {
    if x < 0 {
        return -x
    }
    return x
}

type hp struct{ sort.IntSlice }

func (h hp) Less(i, j int) bool { return h.IntSlice[i] > h.IntSlice[j] }
func (h *hp) Push(v any)        { h.IntSlice = append(h.IntSlice, v.(int)) }
func (h *hp) Pop() any {
    a := h.IntSlice
    v := a[len(a)-1]
    h.IntSlice = a[:len(a)-1]
    return v
}
func (h *hp) push(v int) { heap.Push(h, v) }
func (h *hp) pop() int   { return heap.Pop(h).(int) }

function resultsArray(queries: number[][], k: number): number[] {
    const pq = new MaxPriorityQueue();
    const ans: number[] = [];
    for (const [x, y] of queries) {
        pq.enqueue(Math.abs(x) + Math.abs(y));
        if (pq.size() > k) {
            pq.dequeue();
        }
        ans.push(pq.size() === k ? pq.front().element : -1);
    }
    return ans;
}

3275. K-th Nearest Obstacle Queries

Description

Solutions

Solution 1: Priority Queue (Max-Heap)

Comments