CareerCup

We can just calculate the distance of points from origin and save them in an array. Now you can use Quickselect on the array to find the k closest points.

Looking for interview experience sharing and coaching?

Visit AONECODE.COM for ONE-TO-ONE private lessons by FB, Google and Uber engineers!

SYSTEM DESIGN Courses (highly recommended for candidates of FB, LinkedIn, Amazon, Google & Uber etc.)
ALGORITHMS (conquer DP, Greedy, Graph, Advanced Algorithms & Clean Coding),
latest interview questions sorted by companies,
mock interviews.

Our students got hired from G, U, FB, Amazon, LinkedIn, MS and other top-tier companies after weeks of training.

Email us aonecoding@gmail.com with any questions. Thanks!


SOLUTION:

public class KClosestPoints {

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

    public List<Point> findKClosest(Point[] p, int k) {
        PriorityQueue<Point> pq = new PriorityQueue<>(k, new Comparator<Point>() {
            @Override
            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 (int i = 0; i < p.length; i++) {
            if (i < k)
                pq.offer(p[i]);
            else {
                Point temp = pq.peek();
                if ((p[i].x * p[i].x + p[i].y * p[i].y) - (temp.x * temp.x + temp.y * temp.y) < 0) {
                    pq.poll();
                    pq.offer(p[i]);
                }
            }
        }

        List<Point> x = new ArrayList<>();
        while (!pq.isEmpty())
            x.add(pq.poll());

        return x;
    }
}

a

Build Max heap with first k elements from array. Value of heap elements is distance from 0,0. Then for remaining elements in array, compare with top element of heap.

Quickselect. No extra space needed. O(n) expected time, O(n*n) worst case time

def nSmallest(ps, n, key):
    beg = 0
    end = len(ps)
    while beg != end:
        first = beg
        last = end-1
        piv = last
        pivk = key(piv)
        while first < last:
            if key(first) < pivk:
                first += 1
            else:
                last -= 1
                ps[first],ps[last] = ps[last], ps[first]
        ps[first],ps[piv] = ps[piv], ps[first]
        if first < n-1:
            beg = first + 1
        elif first > n-1:
            end = first
        else:
            return ps[:first+1]
    return ps[:beg]

def nNearest(ps, n):
    def key(i):
        (px, py) = ps[i]
        return px*px + py*py
    return nSmallest(ps, n, key)

O(n) extra space if mutating the input is not allowed:

def nSmallest(ps, n, key):
    r = []
    while ps:
        piv = ps[-1]
        pivk = key(piv)
        lh, rh = [], [piv]
        for t in ps[:-1]:
            (lh if key(t) < pivk else rh).append(t)
        if len(lh) < n:
            r += lh
            ps = rh
            n -= len(lh)
        elif len(lh) > n:
            ps = lh
        else:
            return r+lh
    return r

A

points = [(1,3), (3,2), (4,5), (2,4), (2,3)]

k = 3

from math import sqrt

def euclidian_dist(p):
    return sqrt(p[0]**2+p[1]**2)


def swap(a, i, j):
    tmp = a[i]
    a[i] = a[j]
    a[j] = tmp

def partition(a, pos, start, end):
    new_pos = pos
    for i in range(start, end+1):
        if euclidian_dist(a[new_pos]) > euclidian_dist(a[i]):
            swap(a, new_pos, i)
            print(a)
            new_pos = i
    return new_pos

def quicksort(a, start, end):
    if start < end:
        p = partition(a, start, start+1, end)
        print(a)
        quicksort(a, start, p-1)
        quicksort(a, p+1, end)

def knn(a,k):
    quicksort(a, 0, len(a)-1)        
    return a[:k]

print('result:',knn(points, 3))

public class KClosestPoints {

    private static class Point{
        int x, y;

        public Point(int x, int y){
            this.x = x;
            this.y = y;
        }
        public String toString(){
            return x + " ," + y;
        }
    }

    public static void main(String[] args) {
        findKClosest(Arrays.asList(new Point[]{new Point(1,3), new Point(3, 2), new Point(4, 5), new Point(2,4), new Point(2,3)}), 3).forEach(System.out::println);
    }

    public static List<Point> findKClosest(List<Point> points, int k){
        Comparator<Point> comp = (a, b) -> (a.y *a.y + a.x * a.x) - (b.y * b.y + b.x *b.x);
        Collections.sort(points, comp);
        List<Point> result = new ArrayList<>();
        for(Point p : points){
            if(result.size() < k){
                result.add(p);
            } else {
                break;
            }
        }
        return result;
    }
}

import math

points = [(2, 4), (4, 5), (2, 6), (5, 1), (5, 2), (6, 4), (1, 1), (2, 3)]



def closest_points(pts, k):
    return sorted(pts, key=lambda kv: math.sqrt(kv[0]**2 + kv[1]**2))[:k]


print(closest_points(points, 3))