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);
        }
    }

}