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

@@ -1,6 +1,6 @@
 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):
+    for reflected_position, 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):
@@ -10,7 +10,7 @@ def solution(dimentions, your_position, trainer_position, distance):
             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):
+def generate_reflected_positions(dimentions, your_position, trainer_position, distance):
     room_width, room_height = dimentions
     x_range = distance // room_width + 1
     y_range = distance // room_height + 1
@@ -44,8 +44,12 @@ def abs_gcd(x, y):
     return abs(m)
 
 def reduce_direction(direction):
-    if direction[0] == 0 or direction[1] == 0:
-        return direction
+    if direction[0] == 0 and direction[1] == 0:
+        return 0, 0
+    elif direction[0] == 0:
+        return 0, direction[1] / abs(direction[1])
+    elif direction[1] == 0:
+        return direction[0] / abs(direction[0]), 0
     gcd = abs_gcd(direction[0], direction[1])
     return direction[0]/gcd, direction[1]/gcd