nicer code

1 files changed, 16 insertions, 21 deletions

diff --git a/capturer.py b/capturer.py
index a6bb990..92f85b7 100644
--- a/capturer.py
+++ b/capturer.py
@@ -8,35 +8,30 @@ from robots.utils import add_spawns
 
 class CaptureSpawns(object):
     def __init__(self, variance=0.1):
-        self.distances = None
         self.iterations = 10
         self.variance = variance
 
     def __call__(self, whoami, state):
         my_robots = state.robots_by_player[whoami]
 
-        self.distances = [
-            [
-                {
-                    City.NORMAL: (
-                        float('inf')
-#                        if self.distances is None or self.distances[y][x] is None
-#                        else self.distances[y][x] + self.iterations // 2
-                    ),
-                    City.FACTORY: (
-                        0
-                        if state.allegiances.get((x, y)) != whoami else
-                        float('inf')
-                    ),
-                    City.GHOST: None,
-                }[city]
-                for x, city in enumerate(row)
-            ]
-            for y, row in enumerate(state.cities)
-        ]
+        # Create a distance matrix.
+        distances = []
+        for y, row in enumerate(state.cities):
+            output = []
+            for x, city in enumerate(row):
+                d = float('inf')
+                if city == City.FACTORY:
+                    if state.allegiances.get((x, y)) != whoami:
+                        d = 0
+                elif city == City.GHOST:
+                    d = None
+                output.append(d)
+            distances.append(output)
 
-        predecessors = distance_relaxer(self.distances, None)
+        # Find the shortest path to a target from each cell.
+        predecessors = distance_relaxer(distances)
 
+        # Direct the robots to follow those paths.
         results = []
         for x, y, energy in my_robots:
             if random.random() < self.variance: