Optimize dijkstra mapping for volume calculation.

[plomrogue2] / plomrogue / mapping.py

diff --git a/plomrogue/mapping.py b/plomrogue/mapping.py
index 58809ef8a2764cf7e3212e59687ae9e938013f86..01c10e893c872322c44e7dab1c8bc571d270bf16 100644 (file)
--- a/plomrogue/mapping.py
+++ b/plomrogue/mapping.py
@@ -20,6 +20,7 @@ class MapGeometry():
 
     def __init__(self, size):
         self.size = size
+        self.neighbors_i = {}
 
     def get_directions(self):
         directions = []
@@ -34,6 +35,21 @@ class MapGeometry():
             neighbors[direction] = self.move(pos, direction)
         return neighbors
 
+    def get_neighbors_i(self, i):
+        if i in self.neighbors_i:
+            return self.neighbors_i[i]
+        pos = YX(i // self.size.x, i % self.size.x)
+        neighbors_pos = self.get_neighbors(pos)
+        neighbors_i = {}
+        for direction in neighbors_pos:
+            pos = neighbors_pos[direction]
+            if pos is None:
+                neighbors_i[direction] = None
+            else:
+                neighbors_i[direction] = pos.y * self.size.x + pos.x
+        self.neighbors_i[i] = neighbors_i
+        return self.neighbors_i[i]
+
     def move(self, start_pos, direction):
         mover = getattr(self, 'move_' + direction)
         target = mover(start_pos)