FB面经 Prepare: K closest point to the origin

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

 1 package fbPractise;
 2 
 3 import java.util.*;
 4 
 5 class Coordinate {
 6     int x;
 7     int y;
 8     public Coordinate(int x, int y) {
 9         this.x = x;
10         this.y = y;
11     }
12 }
13 
14 public class Kclosest {
15     
16     public static List<Coordinate> findK(List<Coordinate> input, int k) {
17         HashMap<Coordinate, Integer> map = new HashMap<Coordinate, Integer>();
18         int longest = 0;
19         for (Coordinate each : input) {
20             int distance = cal(each);
21             map.put(each, distance);
22             longest = Math.max(longest, distance);
23         }
24         
25         List<Coordinate>[] arr = new ArrayList[longest + 1];
26         for (Coordinate each : map.keySet()) {
27             int dis = map.get(each);
28             if (arr[dis] == null)
29                 arr[dis] = new ArrayList<Coordinate>();
30             arr[dis].add(each);
31         }
32         
33         List<Coordinate> res = new ArrayList<Coordinate>();
34         for (int i=0; i<arr.length-1 && res.size()<k; i++) {
35             if (arr[i] != null)
36                 res.addAll(arr[i]);
37         }
38         return res;
39     }
40     
41     public static int cal(Coordinate a) {
42         return a.x * a.x + a.y * a.y;
43     }
44     
45 
46     /**
47      * @param args
48      */
49     public static void main(String[] args) {
50         // TODO Auto-generated method stub
51         Coordinate c1 = new Coordinate(1,2);
52         Coordinate c2 = new Coordinate(1,3);
53         Coordinate c3 = new Coordinate(2,5);
54         List<Coordinate> list = new ArrayList<Coordinate>();
55         list.add(c1);
56         list.add(c2);
57         list.add(c3);
58         List<Coordinate> res = findK(list, 2);
59         for (Coordinate each : res) {
60             System.out.print(each.x);
61             System.out.println(each.y);
62         }
63     }
64 
65 }

 1 package fbOnsite;
 2 
 3 import java.util.*;
 4 
 5 class Coordinate {
 6     int x;
 7     int y;
 8     public Coordinate(int x, int y) {
 9         this.x = x;
10         this.y = y;
11     }
12 }
13 
14 public class ClosestKPoints {
15     
16     public static List<Coordinate> findK(List<Coordinate> input, int k, Coordinate target) {
17         HashMap<Coordinate, Integer> map = new HashMap<Coordinate, Integer>();
18         for (Coordinate each : input) {
19             int distance = cal(each, target);
20             map.put(each, distance);
21         }
22         List<Coordinate> res = help(input, 0, input.size()-1, k, map);
23         return res;
24     }
25     
26     public static List<Coordinate> help(List<Coordinate> input, int start, int end, int k, HashMap<Coordinate, Integer> map) {
27         List<Coordinate> res = new ArrayList<Coordinate>();
28         int l = start, r = end;
29         int pivot = r;
30         while (l < r) {
31             while (l<r && map.get(input.get(l))<map.get(input.get(pivot))) l++;
32             while (l<r && map.get(input.get(r))>=map.get(input.get(pivot))) r--;
33             if (l >= r) break;
34             swap(input, l, r);
35         }
36         swap(input, l, pivot);
37         if (l+1 == k) {
38             for (int i=0; i<=l; i++) {
39                 res.add(input.get(i));
40             }
41             return res;
42         }
43         else if (l+1 < k) {
44             return help(input, l+1, end, k, map);
45         }
46         else return help(input, start, l-1, k, map);
47     }
48     
49     public static int cal(Coordinate a, Coordinate target) {
50         return (a.x-target.x)*(a.x-target.x) + (a.y-target.y)*(a.y-target.y);
51     }
52     
53     public static void swap(List<Coordinate> input, int l, int r) {
54         Coordinate temp = input.get(l);
55         input.set(l, input.get(r));
56         input.set(r, temp);
57     }
58     
59 
60     /**
61      * @param args
62      */
63     public static void main(String[] args) {
64         // TODO Auto-generated method stub
65         Coordinate c1 = new Coordinate(1,2);
66         Coordinate c2 = new Coordinate(1,3);
67         Coordinate c3 = new Coordinate(2,5);
68         List<Coordinate> list = new ArrayList<Coordinate>();
69         list.add(c1);
70         list.add(c2);
71         list.add(c3);
72         List<Coordinate> res = findK(list, 2, new Coordinate(2,6));
73         for (Coordinate each : res) {
74             System.out.print(each.x);
75             System.out.println(each.y);
76         }
77     }
78 
79 }

 1 package fbOnsite;
 2 import java.util.*;
 3 public class Kpoints {
 4     class Point {
 5         int x;
 6         int y;
 7         public Point(int x, int y) {
 8             this.x = x;
 9             this.y = y;
10         }
11     }
12     
13     public List<Point> KClosest(List<Point> input, int k) {
14         List<Point> res = new LinkedList<Point>();
15         PriorityQueue<Point> pq = new PriorityQueue<Point>(10, new Comparator<Point>() {
16             public int compare(Point a, Point b) {
17                 return (b.x * b.x + b.y * b.y) - (a.x * a.x + a.y * a.y);
18             }            
19         });
20         
21         for (Point each : input) {
22             pq.offer(each);
23             if (pq.size() > k) pq.poll();
24         }
25         while (!pq.isEmpty()) {
26             res.add(0, pq.poll());
27         }
28         return res;
29     }
30 
31     /**
32      * @param args
33      */
34     public static void main(String[] args) {
35         // TODO Auto-generated method stub
36         Kpoints sol = new Kpoints();
37         List<Point> input = new ArrayList<Point>();
38         input.add(sol.new Point(1,2));
39         input.add(sol.new Point(3,5));
40         input.add(sol.new Point(2,3));
41         input.add(sol.new Point(-1,-7));
42         input.add(sol.new Point(-1,-2));
43         List<Point> res = sol.KClosest(input, 3);
44         for (Point each : res) {
45             System.out.print(each.x);
46             System.out.println(each.y);
47         }
48     }
49 
50 }