973. K Closest Points to Origin

Approach 1: Heap

• Time: $O(n\log K)$
• Space: $O(K)$

C++JavaPython

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 auto& 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]

Approach 2: Quick Select

• Time: $O(n) \to O(n^2)$
• Space: $O(K)$

C++JavaPython

class Solution {
 public:
  vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
    quickSelect(points, 0, points.size() - 1, K);
    return {begin(points), begin(points) + K};
  };

 private:
  void quickSelect(vector<vector<int>>& points, int l, int r, int K) {
    const vector<int> pivot = points[r];

    int nextSwapped = l;
    for (int i = l; i < r; ++i)
      if (squareDist(points[i]) <= squareDist(pivot))
        swap(points[nextSwapped++], points[i]);
    swap(points[nextSwapped], points[r]);

    const int count = nextSwapped - l + 1;  // # of points <= pivot
    if (count == K)
      return;
    if (count > K)
      quickSelect(points, l, nextSwapped - 1, K);
    else
      quickSelect(points, nextSwapped + 1, r, K - count);
  }

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

class Solution {
  public int[][] kClosest(int[][] points, int K) {
    quickSelect(points, 0, points.length - 1, K);
    return Arrays.copyOfRange(points, 0, K);
  }

  private void quickSelect(int[][] points, int l, int r, int K) {
    final int[] pivot = points[r];

    int nextSwapped = l;
    for (int i = l; i < r; ++i)
      if (squareDist(points[i]) <= squareDist(pivot))
        swap(points, nextSwapped++, i);
    swap(points, nextSwapped, r);

    final int count = nextSwapped - l + 1; // # of points <= pivot
    if (count == K)
      return;
    if (count > K)
      quickSelect(points, l, nextSwapped - 1, K);
    else
      quickSelect(points, nextSwapped + 1, r, K - count);
  }

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

  private void swap(int[][] points, int i, int j) {
    final int[] temp = points[i];
    points[i] = points[j];
    points[j] = temp;
  }
}

class Solution:
  def kClosest(self, points: List[List[int]], K: int) -> List[List[int]]:
    def squareDist(p: List[int]) -> int:
      return p[0] * p[0] + p[1] * p[1]

    def quickSelect(l: int, r: int, k: int) -> None:
      pivot = points[r]

      nextSwapped = l
      for i in range(l, r):
        if squareDist(points[i]) <= squareDist(pivot):
          points[nextSwapped], points[i] = points[i], points[nextSwapped]
          nextSwapped += 1
      points[nextSwapped], points[r] = points[r], points[nextSwapped]

      count = nextSwapped - l + 1  # of points <= pivot
      if count == k:
        return
      if count > k:
        quickSelect(l, nextSwapped - 1, k)
      else:
        quickSelect(nextSwapped + 1, r, k - count)

    quickSelect(0, len(points) - 1, K)
    return points[0:K]

Approach 3: Quick Select with random pivot

• Time: $O(n)$ (average)
• Space: $O(K)$

C++Java

class Solution {
 public:
  vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
    quickSelect(points, 0, points.size() - 1, K);
    return {begin(points), begin(points) + K};
  };

 private:
  void quickSelect(vector<vector<int>>& points, int l, int r, int K) {
    const int randIndex = rand() % (r - l + 1) + l;
    swap(points[randIndex], points[r]);
    const vector<int> pivot = points[r];

    int nextSwapped = l;
    for (int i = l; i < r; ++i)
      if (squareDist(points[i]) <= squareDist(pivot))
        swap(points[nextSwapped++], points[i]);
    swap(points[nextSwapped], points[r]);

    const int count = nextSwapped - l + 1;  // # of points <= pivot
    if (count == K)
      return;
    if (count > K)
      quickSelect(points, l, nextSwapped - 1, K);
    else
      quickSelect(points, nextSwapped + 1, r, K - count);
  }

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

class Solution {
  public int[][] kClosest(int[][] points, int K) {
    quickSelect(points, 0, points.length - 1, K);
    return Arrays.copyOfRange(points, 0, K);
  }

  private void quickSelect(int[][] points, int l, int r, int K) {
    final int randIndex = new Random().nextInt(r - l + 1) + l;
    swap(points, randIndex, r);
    final int[] pivot = points[r];

    int nextSwapped = l;
    for (int i = l; i < r; ++i)
      if (squareDist(points[i]) <= squareDist(pivot))
        swap(points, nextSwapped++, i);
    swap(points, nextSwapped, r);

    final int count = nextSwapped - l + 1; // # of points <= pivot
    if (count == K)
      return;
    if (count > K)
      quickSelect(points, l, nextSwapped - 1, K);
    else
      quickSelect(points, nextSwapped + 1, r, K - count);
  }

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

  private void swap(int[][] points, int i, int j) {
    final int[] temp = points[i];
    points[i] = points[j];
    points[j] = temp;
  }
}