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 representing your location coordinates

post_offices: a list representing 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

1from heapq import heapify, heappop
2from math import dist
3
4def nearest_post_offices(you: list[int], offices: list[list[int]], k: int) -> list[list[int]]:
5    h = [(dist(you, o), o) for o in offices]
6    heapify(h)
7    return [heappop(h)[1] for _ in range(k)]
8
9if __name__ == "__main__":
10    you = [int(x) for x in input().split()]
11    offices = [[int(x) for x in input().split()] for _ in range(int(input()))]
12    k = int(input())
13    res = nearest_post_offices(you, offices, k)
14    for row in res:
15        print(" ".join(map(str, row)))
16