import math def solution(dimentions, your_position, trainer_position, distance): direction_target_map = {} for position, direction, is_trainer in generate_reflected_positions(dimentions, your_position, trainer_position, distance): if direction in direction_target_map: prev_is_trainer, prev_position = direction_target_map[direction] if abs(position[0] - your_position[0]) < abs(prev_position[0] - your_position[0]): direction_target_map[direction] = (is_trainer, position) else: direction_target_map[direction] = (is_trainer, position) return len([0 for d, (is_trainer, position) in direction_target_map.items() if is_trainer and line_length(your_position, position) <= distance]) def generate_room_id(dimentions, your_position, distance): memo = [] room_width, room_height = dimentions x_range = distance / room_width + 1 y_range = distance / room_height + 1 for room_x in range(x_range): for room_y in range(y_range): if room_x != 0 and room_y != 0: if line_length(dimentions, (room_x*dimentions[0], room_y*dimentions[1])) > distance: continue memo += [(room_x, room_y)] yield room_x, room_y for room_x, room_y in memo: yield - room_x - 1, room_y yield room_x, - room_y - 1 yield - room_x - 1, - room_y - 1 def generate_reflected_positions(dimentions, your_position, trainer_position, distance): memo = {} 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, 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 your_direction = calc_direction(your_position, reflected_your_position) if your_direction not in memo: yield reflected_your_position, your_direction, False reflected_trainer_position = x_base + trainer_x_offset, y_base + trainer_y_offset trainer_direction = calc_direction(your_position, reflected_trainer_position) if trainer_direction not in memo: yield reflected_trainer_position, trainer_direction, True memo[your_direction] = 0 memo[trainer_direction] = 0 def abs_gcd(x, y): m, n = abs(x), abs(y) while n != 0: t = m % n m, n = n, t return m def reduce_direction(direction): if direction[0] == 0 and direction[1] == 0: return 0, 0 elif direction[0] == 0: return 0, (1 if direction[1] > 0 else -1) elif direction[1] == 0: return (1 if direction[0] > 0 else -1), 0 gcd = abs_gcd(direction[0], direction[1]) return direction[0]/gcd, direction[1]/gcd 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], 1), 1), (([3, 2], [1, 1], [2, 1], 2), 1), #(([3, 2], [1, 1], [2, 1], 10000), 1), ] for i, o in tests: result = solution(*i) print(i, result == o, result, o)