extra/bringing-a-gun-to-a-trainer-fight/solution.py

@@ -1,21 +1,26 @@
 def solution(dimentions, your_position, trainer_position, distance):
     direction_target_map = {}
-    for reflected_position, is_trainer in generate_reflected_positions(dimentions, your_position, trainer_position, distance):
+    for reflected_position, direction, length, is_trainer in generate_reflected_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]:
+            if length < direction_target_map[direction][0]:
-                direction_target_map[direction] = (travel, is_trainer)
+                direction_target_map[direction] = (length, is_trainer)
         else:
-            direction_target_map[direction] = (travel, is_trainer)
+            direction_target_map[direction] = (length, is_trainer)
     return len([k for k, (d, is_trainer) in direction_target_map.items() if is_trainer])
 
+def generate_room_id(max):
+    for room_id in range(max):
+        yield room_id
+    for room_id in range(max):
+        yield - room_id - 1
+
 def generate_reflected_positions(dimentions, your_position, trainer_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, x_range):
+    for room_x in generate_room_id(x_range):
-        for room_y in range(-y_range, y_range):
+        for room_y in generate_room_id(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
@@ -29,12 +34,18 @@ def generate_reflected_positions(dimentions, your_position, trainer_position, di
                 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_your_direction = calc_direction(your_position, reflected_your_position)
-            reflected_trainer_distance = travel_distance(your_position, reflected_trainer_position)
+            reflected_trainer_direction = calc_direction(your_position, reflected_trainer_position)
-            if reflected_your_distance <= distance:
+            if reflected_your_direction not in memo:
-                yield reflected_your_position, False
+                reflected_your_distance = line_length(your_position, reflected_your_position)
-            if reflected_trainer_distance <= distance:
+                if reflected_your_distance <= distance:
-                yield reflected_trainer_position, True
+                    yield reflected_your_position, reflected_your_direction, reflected_your_distance, False
+            if reflected_trainer_direction not in memo:
+                reflected_trainer_distance = line_length(your_position, reflected_trainer_position)
+                if reflected_trainer_distance <= distance:
+                    yield reflected_trainer_position, reflected_trainer_direction, reflected_trainer_distance, True
+            memo[reflected_your_direction] = 0
+            memo[reflected_trainer_direction] = 0
 
 def abs_gcd(x, y):
     m, n = abs(x), abs(y)
@@ -43,6 +54,9 @@ def abs_gcd(x, y):
         m, n = n, t
     return abs(m)
 
+def calc_direction(origin, target):
+    return reduce_direction((target[0] - origin[0], target[1] - origin[1]))
+
 def reduce_direction(direction):
     if direction[0] == 0 and direction[1] == 0:
         return 0, 0
@@ -53,7 +67,7 @@ def reduce_direction(direction):
     gcd = abs_gcd(direction[0], direction[1])
     return direction[0]/gcd, direction[1]/gcd
 
-def travel_distance(origin, target):
+def line_length(origin, target):
     return ((target[0] - origin[0])**2 + (target[1] - origin[1])**2) ** 0.5
 
 tests = [