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

@@ -6,7 +6,7 @@ def solution(dimentions, your_position, trainer_position, distance):
         direction = calc_direction(your_position, position)
         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]):
+            if abs(position[0] - your_position[0]) < abs(prev_position[0] - your_position[0]) or abs(position[1] - your_position[1]) < abs(prev_position[1] - your_position[1]):
                 nearest_target[direction] = (position, is_trainer)
         else:
             nearest_target[direction] = (position, is_trainer)
@@ -17,9 +17,17 @@ def generate_room_id(dimentions, your_position, distance):
     room_height = dimentions[1]
     x_range = distance / room_width + 1
     y_range = distance / room_height + 1
-    for room_x in range(-x_range, x_range):
-        for room_y in range(-y_range, y_range):
-            yield room_x, room_y
+    boundary = {0: y_range + 1}
+    for x in range(1, x_range):
+        while line_length(your_position, (x*room_width, y_range*room_height)) > distance and y_range > 0:
+            y_range -= 1
+        boundary[x] = y_range + 1
+    for x in range(x_range):
+        for y in range(boundary[x]):
+            yield x, y
+            yield x, - y - 1
+            yield - x - 1, y
+            yield - x - 1, - y - 1
 
 def generate_positions(dimentions, your_position, trainer_position, distance):
     room_width, room_height = dimentions