LeetCode-in-Java

3275. K-th Nearest Obstacle Queries

Medium

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:

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

Return an integer array results where results[i] denotes the kth 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:

Example 2:

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

Output: [10,8,6]

Explanation:

Constraints:

Solution

public class Solution {
    public int[] resultsArray(int[][] queries, int k) {
        final int len = queries.length;
        int[] results = new int[len];
        int[] heap = new int[k];
        for (int i = 0; i < k && i < len; i++) {
            int[] query = queries[i];
            heap[i] = Math.abs(query[0]) + Math.abs(query[1]);
            results[i] = -1;
        }
        if (k <= len) {
            buildMaxHeap(heap, k);
            results[k - 1] = heap[0];
        }
        for (int i = k; i < len; i++) {
            int[] query = queries[i];
            int dist = Math.abs(query[0]) + Math.abs(query[1]);
            if (dist < heap[0]) {
                heap[0] = dist;
                heapify(heap, 0, k);
            }
            results[i] = heap[0];
        }
        return results;
    }

    private void buildMaxHeap(int[] heap, int size) {
        for (int i = size / 2 - 1; i >= 0; i--) {
            heapify(heap, i, size);
        }
    }

    private void heapify(int[] heap, int index, int size) {
        int root = heap[index];
        final int left = 2 * index + 1;
        final int right = 2 * index + 2;
        if (right < size && root < heap[right] && heap[left] < heap[right]) {
            heap[index] = heap[right];
            heap[right] = root;
            heapify(heap, right, size);
        } else if (left < size && root < heap[left]) {
            heap[index] = heap[left];
            heap[left] = root;
            heapify(heap, left, size);
        }
    }
}