LeetCode 3275 - K-th Nearest Obstacle Queries

Difficulty: 🟡 Medium
Topics: Array, Heap (Priority Queue)

Solution

Problem Understanding

The problem requires tracking obstacles on an infinite 2D plane and answering, after each obstacle is placed, the distance to the k-th nearest obstacle from the origin (0, 0) based on Manhattan distance, defined as |x| + |y|. The input consists of an integer k and a list of queries, where each query specifies the coordinates of a new obstacle. After each query, we must determine the k-th smallest distance among all current obstacles. If fewer than k obstacles exist at that point, the result is -1.

Key points include that all queries are unique, obstacles are never removed, and the input size can be large: up to 2 * 10^5 queries and up to 10^5 for k. This means naive methods that repeatedly sort all distances after each query will be too slow.

Edge cases include scenarios where fewer than k obstacles exist, large coordinates leading to large distances, and k = 1, which corresponds to tracking the nearest obstacle.

Approaches

A brute-force approach would be to maintain a list of all distances, append the new distance on each query, and sort the list to extract the k-th smallest distance. While simple, sorting the list on every query is inefficient: each sort takes O(n log n) and with n queries, the total time is O(n^2 log n) - far too slow for n = 2 * 10^5.

A more optimal approach leverages a max-heap of size k to track the k smallest distances dynamically. For each new distance, if the heap has fewer than k elements, we add it. If it has exactly k elements and the new distance is smaller than the current maximum, we replace the maximum. This ensures the heap always contains the k smallest distances, and the root of the max-heap is the k-th smallest distance. Heap operations are O(log k), giving a total time complexity of O(n log k), which is feasible given the constraints.

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(n^2 log n)	O(n)	Sort distances at each query
Max-Heap Optimal	O(n log k)	O(k)	Maintain max-heap of size k

Algorithm Walkthrough

Initialize an empty max-heap to store up to k distances.
Initialize an empty results array.
For each query (x, y):
Compute the Manhattan distance d = |x| + |y|.
If the heap contains fewer than k elements, push d onto the heap.
Otherwise, compare d with the maximum element of the heap (the root). If d is smaller, replace the root with d.
If the heap size is less than k, append -1 to the results.
Otherwise, append the maximum element in the heap (k-th smallest distance) to results.
Return the results array.

Why it works: By maintaining a max-heap of size k, the root always represents the k-th smallest distance among all inserted distances. We only ever keep the k smallest distances, so extracting the root guarantees correctness. Inserting and replacing elements in the heap preserves the invariant.

Python Solution

from typing import List
import heapq

class Solution:
    def resultsArray(self, queries: List[List[int]], k: int) -> List[int]:
        max_heap = []
        results = []
        
        for x, y in queries:
            dist = abs(x) + abs(y)
            if len(max_heap) < k:
                heapq.heappush(max_heap, -dist)
            else:
                if dist < -max_heap[0]:
                    heapq.heapreplace(max_heap, -dist)
            results.append(-max_heap[0] if len(max_heap) == k else -1)
        
        return results

In this implementation, max_heap stores negative distances because Python's heapq is a min-heap by default. By pushing -dist, the smallest distances appear at the bottom, and the largest among them is always the root. After each query, if the heap size is at least k, we return -max_heap[0] as the k-th smallest distance; otherwise, we return -1.

Go Solution

package main

import (
    "container/heap"
    "math"
)

type MaxHeap []int

func (h MaxHeap) Len() int            { return len(h) }
func (h MaxHeap) Less(i, j int) bool  { return h[i] > h[j] } // max-heap
func (h MaxHeap) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }
func (h *MaxHeap) Push(x any)         { *h = append(*h, x.(int)) }
func (h *MaxHeap) Pop() any {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

func resultsArray(queries [][]int, k int) []int {
    h := &MaxHeap{}
    heap.Init(h)
    results := make([]int, len(queries))

    for i, q := range queries {
        dist := int(math.Abs(float64(q[0])) + math.Abs(float64(q[1])))
        if h.Len() < k {
            heap.Push(h, dist)
        } else if dist < (*h)[0] {
            heap.Pop(h)
            heap.Push(h, dist)
        }

        if h.Len() < k {
            results[i] = -1
        } else {
            results[i] = (*h)[0]
        }
    }

    return results
}

Go implementation differences include defining a custom MaxHeap type because Go's container/heap is min-heap by default. We define Less in reverse to simulate a max-heap. Type casting is necessary when pushing and popping from the heap.

Worked Examples

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

Query	Distance	Heap Contents (max-heap)	k-th Distance	Results
[1,2]	3	[3]	-	-1
[3,4]	7	[7,3]	7	7
[2,3]	5	[7,3,5] → heap keeps 2 smallest → [5,3]	5	5
[-3,0]	3	[5,3,3] → heap keeps 2 smallest → [3,3]	3	3

Example 2: queries = [[5,5],[4,4],[3,3]], k = 1

Query	Distance	Heap Contents	k-th Distance	Results
[5,5]	10	[10]	10	10
[4,4]	8	[10,8] → heap keeps 1 smallest → [8]	8	8
[3,3]	6	[8,6] → heap keeps 1 smallest → [6]	6	6

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n log k)	Each of n queries does up to one heap operation, each `O(log k)`
Space	O(k)	Max-heap stores at most k distances

The reasoning is that heap operations scale logarithmically with the heap size k, and the heap never exceeds size k.

Test Cases

sol = Solution()

# Provided examples
assert sol.resultsArray([[1,2],[3,4],[2,3],[-3,0]], 2) == [-1,7,5,3]  # Example 1
assert sol.resultsArray([[5,5],[4,4],[3,3]], 1) == [10,8,6]          # Example 2

# Edge cases
assert sol.resultsArray([[0,0]], 1) == [0]                           # Single obstacle
assert sol.resultsArray([[0,0]], 2) == [-1]                          # Less than k obstacles
assert sol.resultsArray([[1,0],[0,1]], 2) == [-1,1]                  # k = number of obstacles
assert sol.resultsArray([[10**9,10**9],[-10**9,-10**9]], 1) == [2000000000,2000000000]  # Large coordinates

Test	Why
[[1,2],[3,4],[2,3],[-3,0]], k=2	Validates normal behavior and heap adjustment
[[5,5],[4,4],[3,3]], k=1	Validates tracking nearest obstacle
[[0,0]], k=1	Minimal input, edge of origin
[[0,0]], k=2	Less than k