Amazon Online Assessment (OA) - K Nearest Post Offices

Find the k post offices located closest to you, given your location and a list of locations of all post offices available.

Locations are given in 2D coordinates in [X, Y], where X and Y are integers.

Euclidean distance is applied to find the distance between you and a post office.

Assume your location is [m, n] and the location of a post office is [p, q], the Euclidean distance between the office and you is SquareRoot((m - p) * (m - p) + (n - q) * (n - q)).

K is a positive integer much smaller than the given number of post offices.

Input

The input consists of three arguments:

you: a list represent your location coordinate

post_offices: a list represent post offices location coordinates

k: a positive integer much smaller than the given number of post offices

Output

Return 2D coordinates in [X, Y] of post offices located closest to you

Examples

Example 1:

Input:

you = [0, 0]

post_offices = [[-16, 5], [-1, 2], [4, 3], [10, -2], [0, 3], [-5, -9]]

K = 3

Output: [[-1, 2], [0, 3], [4, 3]]

Try it yourself

1
1
from typing import List
2
2
3
3
def nearestPostOffices(you: List[int], offices: List[List[int]], k: int) -> List[List[int]]:
4
-
    # WRITE YOUR BRILLIANT CODE HERE
4
+
    from heapq import heapify, heappop, nsmallest
5
+
    from math import dist
6
+
7
+
    h = [(dist(you, o), o) for o in offices]
8
+
    heapify(h)
9
+
    return [heappop(h)[1] for _ in range(k)]
10
+
5
11
if __name__ == "__main__":
6
12
    you = [int(y) for y in input().split()]
7
13
    k = int(input());
8
14
    rows = int(input());
Expand 4 lines ...