difficulty: Medium tags:

  • Heap
  • Amazon
  • LinkedIn title: K Closest Points

K Closest Points

Problem

Metadata

Description

Given some points and a point origin in two dimensional space, find k points out of the some points which are nearest to origin. Return these points sorted by distance, if they are same with distance, sorted by x-axis, otherwise sorted by y-axis.

Example

Given points = [[4,6],[4,7],[4,4],[2,5],[1,1]], origin = [0, 0], k = 3 return [[1,1],[2,5],[4,4]]

题解

和普通的字符串及数目比较，此题为距离的比较。

Java

  1. /**
  2.  * Definition for a point.
  3.  * class Point {
  4.  *     int x;
  5.  *     int y;
  6.  *     Point() { x = 0; y = 0; }
  7.  *     Point(int a, int b) { x = a; y = b; }
  8.  * }
  9.  */
  11. public class Solution {
  12.     /**
  13.      * @param points: a list of points
  14.      * @param origin: a point
  15.      * @param k: An integer
  16.      * @return: the k closest points
  17.      */
  18.     public Point[] kClosest(Point[] points, Point origin, int k) {
  19.         // write your code here
  20.         Queue<Point> heap = new PriorityQueue<Point>(new DistanceComparator(origin));
  21.         for (Point point : points) {
  22.             if (heap.size() < k) {
  23.                 heap.offer(point);
  24.             } else {
  25.                 Point peek = heap.peek();
  26.                 if (distance(peek, origin) <= distance(point, origin)) {
  27.                     continue;
  28.                 } else {
  29.                     heap.poll();
  30.                     heap.offer(point);
  31.                 }
  32.             }
  33.         }
  35.         int minK = Math.min(k, heap.size());
  36.         Point[] kClosestPoints = new Point[minK];
  37.         for (int i = 1; i <= minK; i++) {
  38.             kClosestPoints[minK - i] = heap.poll();
  39.         }
  41.         return kClosestPoints;
  42.     }
  44.     public int distance(Point p, Point origin) {
  45.         return (p.x - origin.x) * (p.x - origin.x) + 
  46.                (p.y - origin.y) * (p.y - origin.y);
  47.     }
  49.     class DistanceComparator implements Comparator<Point> {
  50.         private Point origin = null;
  51.         public DistanceComparator(Point origin) {
  52.             this.origin = origin;
  53.         }
  55.         public int compare(Point p1, Point p2) {
  56.             int d1 = distance(p1, origin);
  57.             int d2 = distance(p2, origin);
  58.             if (d1 != d2) {
  59.                 return d2 - d1;
  60.             } else {
  61.                 if (p1.x != p2.x) {
  62.                     return p2.x - p1.x;
  63.                 } else {
  64.                     return p2.y - p1.y;
  65.                 }
  66.             }
  67.         }
  68.     }
  69. }

源码分析

注意 Comparator 的用法和大小根堆的选择即可。

复杂度分析

堆的删除插入操作，最大为 K, 故时间复杂度为 O(n \log k), 空间复杂度为 O(K).