day 18 python cleanup

day18/sol.py

@@ -1,90 +1,63 @@
 import numpy
-lava = numpy.zeros((23,23,23), dtype=numpy.int32)
 
-for line in open("input"):
+N = 23
+
+def solve():
+    lava = numpy.zeros((N,N,N), dtype=numpy.int32)
+    for line in open("input"):
         x,y,z = map(int, line.strip().split(","))
         lava[x,y,z] = 1
 
-water = numpy.zeros((23,23,23), dtype=numpy.int32)
+    lava = numpy.pad(lava, (1,1), constant_values=0)
 
-area = 0
-N = 23
-for x in range(N):
-    for y in range(N):
-        for z in range(N):
-            if lava[x,y,z]:
-                sides = 0
-                def do(x,y,z):
-                    global sides
-                    if 0 <= x < N and 0 <= y < N and 0 <= z < N:
-                        sides += not lava[x,y,z]
-                    else:
-                        sides += 1
-                do(x-1, y, z)
-                do(x+1, y, z)
-                do(x, y-1, z)
-                do(x, y+1, z)
-                do(x, y, z-1)
-                do(x, y, z+1)
-                area += sides
-            else:
-                if 0 in (x,y,z) or N-1 in (x,y,z):
-                    water[x,y,z] = 1
-print(area)
+    # part 1
+    area = count_touching(lava, lava^1)
+    print(area)
 
-def move(w, axis, amount):
-    w = numpy.roll(w, amount, axis=axis)
-    all = slice(None,None)
-    index = [all,all,all]
-    if amount > 1:
-        index[axis] = 0
-    else:
-        index[axis] = w.shape[axis] - 1
-    w[tuple(index)] = 0
-    return w
+    # part 2
+    border = numpy.zeros((N,N,N), dtype=numpy.int32)
+    border = numpy.pad(border, (1,1), constant_values=1)
 
-def flood(a):
+    water = flood(lava, border)
+    #print(water)
+
+    area = count_touching(lava, water)
+    print(area)
+
+
+def count_touching(A, B):
+    """counts the number of places a cell in A touches a cell in B
+    (assumes a 1-cell padding around the edge which is not counted)"""
+    n = 0
+    for x in range(1,N+1):
+        for y in range(1,N+1):
+            for z in range(1,N+1):
+                if A[x,y,z]:
+                    n += B[x-1,y,z]
+                    n += B[x+1,y,z]
+                    n += B[x,y-1,z]
+                    n += B[x,y+1,z]
+                    n += B[x,y,z-1]
+                    n += B[x,y,z+1]
+    return n
+
+def flood(lava, border):
+    def roll(a, axis, amount):
+        return numpy.roll(a, amount, axis=axis)
+
+    a = border
     while True:
-        x = move(a, 0, -1)
-        x |= move(a, 0, +1)
-        x |= move(a, 1, -1)
-        x |= move(a, 1, +1)
-        x |= move(a, 2, -1)
-        x |= move(a, 2, +1)
+        x = roll(a, 0, -1)
+        x |= roll(a, 0, +1)
+        x |= roll(a, 1, -1)
+        x |= roll(a, 1, +1)
+        x |= roll(a, 2, -1)
+        x |= roll(a, 2, +1)
+        x = x & ~border
         x = x & ~lava
         if (a == x).all():
+            a |= border
             return a
         a = x
 
-
-print(move(numpy.ones((3,3,3)), 1, 1))
-print(move(numpy.ones((3,3,3)), 0, -1))
-
-water = flood(water)
-print(water)
-
-area = 0
-N = 23
-for x in range(N):
-    for y in range(N):
-        for z in range(N):
-            if lava[x,y,z]:
-                sides = 0
-                def do(x,y,z):
-                    global sides
-                    if 0 <= x < N and 0 <= y < N and 0 <= z < N:
-                        sides += not lava[x,y,z] and water[x,y,z]
-                    else:
-                        sides += 1
-                do(x-1, y, z)
-                do(x+1, y, z)
-                do(x, y-1, z)
-                do(x, y+1, z)
-                do(x, y, z-1)
-                do(x, y, z+1)
-                area += sides
-            else:
-                if 0 in (x,y,z) or N-1 in (x,y,z):
-                    water[x,y,z] = 1
-print(area)
-
+solve()