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

@@ -1,26 +1,25 @@
 import math
 
 def solution(dimentions, your_position, trainer_position, distance):
-    direction_target_map = {}
+    nearest_target = {}
-    for position, direction, is_trainer in generate_reflected_positions(dimentions, your_position, trainer_position, distance):
+    for position, is_trainer in generate_reflected_positions(dimentions, your_position, trainer_position, distance):
-        if direction in direction_target_map:
+        direction = calc_direction(your_position, position)
-            prev_is_trainer, prev_position = direction_target_map[direction]
+        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]):
-                direction_target_map[direction] = (is_trainer, position)
+                nearest_target[direction] = (position, is_trainer)
         else:
-            direction_target_map[direction] = (is_trainer, position)
+            nearest_target[direction] = (position, is_trainer)
-    return len([0 for d, (is_trainer, position) in direction_target_map.items() if is_trainer and line_length(your_position, position) <= distance])
+    return len([0 for d, (position, is_trainer) in nearest_target.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
+    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):
         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:
@@ -29,30 +28,30 @@ def generate_room_id(dimentions, your_position, distance):
         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
+        x_base = room_x*room_width
-        your_x_offset, your_y_offset = your_position
+        y_base = room_y*room_height
-        trainer_x_offset, trainer_y_offset = trainer_position
+        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):
-            your_x_offset = room_width - your_x_offset
+            reflected_your_x_position = x_base + room_width - your_x_offset
-            trainer_x_offset = room_width - trainer_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):
-            your_y_offset = room_height - your_y_offset
+            reflected_your_y_position = y_base + room_height - your_y_offset
-            trainer_y_offset = room_height - trainer_y_offset
+            reflected_trainer_y_position = y_base + room_height - trainer_y_offset
-        reflected_your_position = x_base + your_x_offset, y_base + your_y_offset
+        else:
-        your_direction = calc_direction(your_position, reflected_your_position)
+            reflected_your_y_position = y_base + your_y_offset
-        if your_direction not in memo:
+            reflected_trainer_y_position = y_base + trainer_y_offset
-            yield reflected_your_position, your_direction, False
+        yield (reflected_your_x_position, reflected_your_y_position), False
-        reflected_trainer_position = x_base + trainer_x_offset, y_base + trainer_y_offset
+        yield (reflected_trainer_x_position, reflected_trainer_y_position), True
-        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)
@@ -62,14 +61,15 @@ def abs_gcd(x, y):
     return m
 
 def reduce_direction(direction):
-    if direction[0] == 0 and direction[1] == 0:
+    x, y = direction
+    if x == 0 and y == 0:
         return 0, 0
-    elif direction[0] == 0:
+    elif x == 0:
-        return 0, (1 if direction[1] > 0 else -1)
+        return 0, (1 if y > 0 else -1)
-    elif direction[1] == 0:
+    elif y == 0:
-        return (1 if direction[0] > 0 else -1), 0
+        return (1 if x > 0 else -1), 0
-    gcd = abs_gcd(direction[0], direction[1])
+    gcd = abs_gcd(x, y)
-    return direction[0]/gcd, direction[1]/gcd
+    return x/gcd, y/gcd
 
 def calc_direction(origin, target):
     return reduce_direction((target[0] - origin[0], target[1] - origin[1]))
@@ -80,9 +80,9 @@ def line_length(origin, target):
 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),
+    (([300,275], [150,150], [185,100], 2), 0),
+    # (([3, 2], [1, 1], [2, 1], 10000), 1),
 ]
 
 for i, o in tests: