LeetCode Solutions

973. K Closest Points to Origin

Time: $O(n\log K)$

Space: $O(K)$

			

class Solution {
 public:
  vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
    vector<vector<int>> ans;
    auto compare = [&](const vector<int>& a, const vector<int>& b) {
      return squareDist(a) < squareDist(b);
    };
    priority_queue<vector<int>, vector<vector<int>>, decltype(compare)> maxHeap(
        compare);

    for (const vector<int>& p : points) {
      maxHeap.push(p);
      if (maxHeap.size() > K)
        maxHeap.pop();
    }

    while (!maxHeap.empty())
      ans.push_back(maxHeap.top()), maxHeap.pop();

    return ans;
  };

 private:
  int squareDist(const vector<int>& p) {
    return p[0] * p[0] + p[1] * p[1];
  }
};
			

class Solution {
  public int[][] kClosest(int[][] points, int K) {
    int[][] ans = new int[K][2];
    PriorityQueue<int[]> maxHeap = new PriorityQueue<>((a, b) -> squareDist(b) - squareDist(a));

    for (int[] p : points) {
      maxHeap.offer(p);
      if (maxHeap.size() > K)
        maxHeap.poll();
    }

    int i = K;
    while (!maxHeap.isEmpty())
      ans[--i] = maxHeap.poll();

    return ans;
  }

  private int squareDist(int[] p) {
    return p[0] * p[0] + p[1] * p[1];
  }
}
			

class Solution:
  def kClosest(self, points: List[List[int]], K: int) -> List[List[int]]:
    maxHeap = []

    for x, y in points:
      heapq.heappush(maxHeap, (- x * x - y * y, [x, y]))
      if len(maxHeap) > K:
        heapq.heappop(maxHeap)

    return [pair[1] for pair in maxHeap]