nearest_stop

MAX_DIST = 20


def calc_dist(xy1, xy2):
    return ((xy1[0] - xy2[0]) ** 2 + (xy1[1] - xy2[1]) ** 2) ** (1 / 2)


def nearest_stop(metro_data, bus_data):
    square_set = set()
    for x in range(-20, 21):
        for y in range(-20, 21):
            if calc_dist((x, y), (0, 0)) <= MAX_DIST:
                square_set.add((x, y))
    max_count = i_max = -1
    for i, (metro_x, metro_y) in enumerate(metro_data):
        count = 0
        for bus_x, bus_y in bus_data:
            if (bus_x - metro_x, bus_y - metro_y) in square_set:
                count += 1
        if count > max_count:
            max_count = count
            i_max = i
    return i_max + 1


a = [[-1, 0], [1, 0], [2, 5]]
b = [[10, 0], [20, 0], [22, 5]]
assert(nearest_stop(a, b) == 3)

a = [[-1, 0], [1, 0], [0, 5]]
b = [[10, 0], [20, 0], [20, 5]]
assert(nearest_stop(a, b) == 2)

a_str = '-87 -73 -68 -71 31 -27 50 -8 44 -7 -33 -106 -74 -50 -30 -37 -42 -92 -78 -44'
b_str = '-106 -55 -37 -46 -91 -40 3 -56 -22 -58 -85 -49 -111 -74 43 -17 60 -3 -81 -73 -40 -40 21 -30 -91 -35 -31 -62 -51 -103'
a = list(map(int, a_str.split()))
a = [[x, y] for x, y in zip(a[::2], a[1::2])]
b = list(map(int, b_str.split()))
b = [[x, y] for x, y in zip(b[::2], b[1::2])]
assert(nearest_stop(a, b) == 10)