FB面经 Prepare: K closest point to the origin

Give n points on 2-D plane, find the K closest points to origin

Based on bucket sort:

 package fbPractise;

 import java.util.*;

 class Coordinate {
     int x;
     int y;
     public Coordinate(int x, int y) {
         this.x = x;
         this.y = y;
     }
 }

 public class Kclosest {

     public static List<Coordinate> findK(List<Coordinate> input, int k) {
         HashMap<Coordinate, Integer> map = new HashMap<Coordinate, Integer>();
         int longest = 0;
         for (Coordinate each : input) {
             int distance = cal(each);
             map.put(each, distance);
             longest = Math.max(longest, distance);
         }

         List<Coordinate>[] arr = new ArrayList[longest + 1];
         for (Coordinate each : map.keySet()) {
             int dis = map.get(each);
             if (arr[dis] == null)
                 arr[dis] = new ArrayList<Coordinate>();
             arr[dis].add(each);
         }

         List<Coordinate> res = new ArrayList<Coordinate>();
         for (int i=0; i<arr.length-1 && res.size()<k; i++) {
             if (arr[i] != null)
                 res.addAll(arr[i]);
         }
         return res;
     }

     public static int cal(Coordinate a) {
         return a.x * a.x + a.y * a.y;
     }

     /**
      * @param args
      */
     public static void main(String[] args) {
         // TODO Auto-generated method stub
         Coordinate c1 = new Coordinate(1,2);
         Coordinate c2 = new Coordinate(1,3);
         Coordinate c3 = new Coordinate(2,5);
         List<Coordinate> list = new ArrayList<Coordinate>();
         list.add(c1);
         list.add(c2);
         list.add(c3);
         List<Coordinate> res = findK(list, 2);
         for (Coordinate each : res) {
             System.out.print(each.x);
             System.out.println(each.y);
         }
     }

 }

Based on Quick Select

 package fbOnsite;

 import java.util.*;

 class Coordinate {
     int x;
     int y;
     public Coordinate(int x, int y) {
         this.x = x;
         this.y = y;
     }
 }

 public class ClosestKPoints {

     public static List<Coordinate> findK(List<Coordinate> input, int k, Coordinate target) {
         HashMap<Coordinate, Integer> map = new HashMap<Coordinate, Integer>();
         for (Coordinate each : input) {
             int distance = cal(each, target);
             map.put(each, distance);
         }
         List<Coordinate> res = help(input, 0, input.size()-1, k, map);
         return res;
     }

     public static List<Coordinate> help(List<Coordinate> input, int start, int end, int k, HashMap<Coordinate, Integer> map) {
         List<Coordinate> res = new ArrayList<Coordinate>();
         int l = start, r = end;
         int pivot = r;
         while (l < r) {
             while (l<r && map.get(input.get(l))<map.get(input.get(pivot))) l++;
             while (l<r && map.get(input.get(r))>=map.get(input.get(pivot))) r--;
             if (l >= r) break;
             swap(input, l, r);
         }
         swap(input, l, pivot);
         if (l+1 == k) {
             for (int i=0; i<=l; i++) {
                 res.add(input.get(i));
             }
             return res;
         }
         else if (l+1 < k) {
             return help(input, l+1, end, k, map);
         }
         else return help(input, start, l-1, k, map);
     }

     public static int cal(Coordinate a, Coordinate target) {
         return (a.x-target.x)*(a.x-target.x) + (a.y-target.y)*(a.y-target.y);
     }

     public static void swap(List<Coordinate> input, int l, int r) {
         Coordinate temp = input.get(l);
         input.set(l, input.get(r));
         input.set(r, temp);
     }

     /**
      * @param args
      */
     public static void main(String[] args) {
         // TODO Auto-generated method stub
         Coordinate c1 = new Coordinate(1,2);
         Coordinate c2 = new Coordinate(1,3);
         Coordinate c3 = new Coordinate(2,5);
         List<Coordinate> list = new ArrayList<Coordinate>();
         list.add(c1);
         list.add(c2);
         list.add(c3);
         List<Coordinate> res = findK(list, 2, new Coordinate(2,6));
         for (Coordinate each : res) {
             System.out.print(each.x);
             System.out.println(each.y);
         }
     }

 }

当然，还有的方法是维护一个size为k的最大堆

 package fbOnsite;
 import java.util.*;
 public class Kpoints {
     class Point {
         int x;
         int y;
         public Point(int x, int y) {
             this.x = x;
             this.y = y;
         }
     }

     public List<Point> KClosest(List<Point> input, int k) {
         List<Point> res = new LinkedList<Point>();
         PriorityQueue<Point> pq = new PriorityQueue<Point>(10, new Comparator<Point>() {
             public int compare(Point a, Point b) {
                 return (b.x * b.x + b.y * b.y) - (a.x * a.x + a.y * a.y);
             }
         });

         for (Point each : input) {
             pq.offer(each);
             if (pq.size() > k) pq.poll();
         }
         while (!pq.isEmpty()) {
             res.add(0, pq.poll());
         }
         return res;
     }

     /**
      * @param args
      */
     public static void main(String[] args) {
         // TODO Auto-generated method stub
         Kpoints sol = new Kpoints();
         List<Point> input = new ArrayList<Point>();
         input.add(sol.new Point(1,2));
         input.add(sol.new Point(3,5));
         input.add(sol.new Point(2,3));
         input.add(sol.new Point(-1,-7));
         input.add(sol.new Point(-1,-2));
         List<Point> res = sol.KClosest(input, 3);
         for (Point each : res) {
             System.out.print(each.x);
             System.out.println(each.y);
         }
     }

 }