import math def solution(dimentions, your_position, trainer_position, distance): nearest_target = {} for position, is_trainer in generate_positions(dimentions, your_position, trainer_position, distance): direction = calc_direction(your_position, position) if direction in nearest_target: prev_position, prev_is_trainer = nearest_target[direction] if abs(position[0] - your_position[0]) < abs(prev_position[0] - your_position[0]): nearest_target[direction] = (position, is_trainer) else: nearest_target[direction] = (position, is_trainer) return sum(1 for (position, is_trainer) in nearest_target.values() if is_trainer and line_length(your_position, position) <= distance) def generate_room_id(dimentions, your_position, distance): room_width = dimentions[0] room_height = dimentions[1] x_range = distance / room_width + 1 y_range = distance / room_height + 1 for room_x in range(-x_range, x_range): for room_y in range(-y_range, y_range): yield room_x, room_y def generate_positions(dimentions, your_position, trainer_position, distance): room_width, room_height = dimentions for room_x, room_y in generate_room_id(dimentions, your_position, distance): x_reflected = room_x % 2 == 1 y_reflected = room_y % 2 == 1 x_base = room_x*room_width y_base = room_y*room_height your_x_offset = your_position[0] your_y_offset = your_position[1] trainer_x_offset = trainer_position[0] trainer_y_offset = trainer_position[1] if (x_reflected): reflected_your_x_position = x_base + room_width - your_x_offset reflected_trainer_x_position = x_base + room_width - trainer_x_offset else: reflected_your_x_position = x_base + your_x_offset reflected_trainer_x_position = x_base + trainer_x_offset if (y_reflected): reflected_your_y_position = y_base + room_height - your_y_offset reflected_trainer_y_position = y_base + room_height - trainer_y_offset else: reflected_your_y_position = y_base + your_y_offset reflected_trainer_y_position = y_base + trainer_y_offset yield (reflected_your_x_position, reflected_your_y_position), False yield (reflected_trainer_x_position, reflected_trainer_y_position), True def gcd(m, n): while n != 0: t = m % n m, n = n, t return m def reduce_direction(direction): x, y = direction if x == 0 and y == 0: return 0, 0 elif x == 0: return 0, (1 if y > 0 else -1) elif y == 0: return (1 if x > 0 else -1), 0 g = gcd(abs(x), abs(y)) return x/g, y/g def calc_direction(origin, target): return reduce_direction((target[0] - origin[0], target[1] - origin[1])) def line_length(origin, target): return math.sqrt((target[0] - origin[0])**2 + (target[1] - origin[1])**2) tests = [ (([3, 2], [1, 1], [2, 1], 4), 7), (([300,275], [150,150], [185,100], 500), 9), (([3, 2], [1, 1], [2, 1], 2), 1), (([300,275], [150,150], [185,100], 2), 0), (([3, 2], [1, 1], [2, 1], 1000), 397845), (([30, 20], [10, 10], [20, 10], 10000), 397845), ] for i, o in tests: result = solution(*i) print(i, result == o, result, o)