def solution(dimentions, your_position, trainer_position, distance): direction_target_map = {} for reflected_position, is_trainer in generate_relected_positions(dimentions, your_position, trainer_position, distance): direction = reduce_direction((reflected_position[0] - your_position[0], reflected_position[1] - your_position[1])) travel = travel_distance(your_position, reflected_position) if (direction in direction_target_map): if travel < direction_target_map[direction][0]: direction_target_map[direction] = (travel, is_trainer) else: direction_target_map[direction] = (travel, is_trainer) return len([k for k, (d, is_trainer) in direction_target_map.items() if is_trainer]) def generate_relected_positions(dimentions, your_position, trainer_position, distance): room_width, room_height = dimentions 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): x_reflected = room_x % 2 == 1 y_reflected = room_y % 2 == 1 x_base, y_base = room_x*room_width, room_y*room_height your_x_offset, your_y_offset = your_position trainer_x_offset, trainer_y_offset = trainer_position if (x_reflected): your_x_offset = room_width - your_x_offset trainer_x_offset = room_width - trainer_x_offset if (y_reflected): your_y_offset = room_height - your_y_offset trainer_y_offset = room_height - trainer_y_offset reflected_your_position = x_base + your_x_offset, y_base + your_y_offset reflected_trainer_position = x_base + trainer_x_offset, y_base + trainer_y_offset reflected_your_distance = travel_distance(your_position, reflected_your_position) reflected_trainer_distance = travel_distance(your_position, reflected_trainer_position) if reflected_your_distance <= distance: yield reflected_your_position, False if reflected_trainer_distance <= distance: yield reflected_trainer_position, True def abs_gcd(x, y): m, n = abs(x), abs(y) while n != 0: t = m % n m, n = n, t return abs(m) def reduce_direction(direction): if direction[0] == 0 or direction[1] == 0: return direction gcd = abs_gcd(direction[0], direction[1]) return direction[0]/gcd, direction[1]/gcd def travel_distance(origin, target): return ((target[0] - origin[0])**2 + (target[1] - origin[1])**2) ** 0.5 tests = [ (([3, 2], [1, 1], [2, 1], 4), 7), (([300,275], [150,150], [185,100], 500), 9) ] for i, o in tests: result = solution(*i) print(i, result == o, result, o)