memo some more variables

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

@@ -1,26 +1,44 @@
+import math
+
 def solution(dimentions, your_position, trainer_position, distance):
     direction_target_map = {}
+    count = 0
     for reflected_position, direction, length, is_trainer in generate_reflected_positions(dimentions, your_position, trainer_position, distance):
-        if (direction in direction_target_map):
+        if direction in direction_target_map:
-            if length < direction_target_map[direction][0]:
+            prev_length, prev_is_trainer = direction_target_map[direction]
+            if length < prev_length:
+                if is_trainer and ~prev_is_trainer:
+                    count += 1
+                elif ~is_trainer and prev_is_trainer:
+                    count -= 1
                 direction_target_map[direction] = (length, is_trainer)
         else:
+            if is_trainer:
+                count += 1
             direction_target_map[direction] = (length, is_trainer)
-    return len([k for k, (d, is_trainer) in direction_target_map.items() if is_trainer])
+    return count
 
-def generate_room_id(max):
+def generate_room_id(dimentions, your_position, distance):
-    for room_id in range(max):
+    memo = []
-        yield room_id
+    room_width, room_height = dimentions
-    for room_id in range(max):
+    x_range = distance / room_width + 1
-        yield - room_id - 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
-    x_range = distance // room_width + 1
+    for room_x, room_y in generate_room_id(dimentions, your_position, distance):
-    y_range = distance // room_height + 1
-    for room_x in generate_room_id(x_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
@@ -33,46 +51,49 @@ def generate_reflected_positions(dimentions, your_position, trainer_position, di
             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
+        your_direction = calc_direction(your_position, reflected_your_position)
-            reflected_your_direction = calc_direction(your_position, reflected_your_position)
+        if your_direction not in memo:
-            reflected_trainer_direction = calc_direction(your_position, reflected_trainer_position)
-            if reflected_your_direction not in memo:
             reflected_your_distance = line_length(your_position, reflected_your_position)
             if reflected_your_distance <= distance:
-                    yield reflected_your_position, reflected_your_direction, reflected_your_distance, False
+                yield reflected_your_position, your_direction, reflected_your_distance, False
-            if reflected_trainer_direction not in memo:
+        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:
             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
+                yield reflected_trainer_position, trainer_direction, reflected_trainer_distance, True
-            memo[reflected_your_direction] = 0
+        memo[your_direction] = 0
-            memo[reflected_trainer_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 abs(m)
+    return 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
     elif direction[0] == 0:
-        return 0, direction[1] / abs(direction[1])
+        return 0, (1 if direction[1] > 0 else -1)
     elif direction[1] == 0:
-        return direction[0] / abs(direction[0]), 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 ((target[0] - origin[0])**2 + (target[1] - origin[1])**2) ** 0.5
+    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)
+    (([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: