aoc_2022_day_18_tf.py

from collections import *
import math

class DSU_Iterative_By_Size:
    """Iterative find, rank by size, and path compression
    https://judge.yosupo.jp/submission/113655"""

    def __init__(self, n):
        self.par = list(range(n))
        self.size = [1] * n

    def find(self, a):
        x = a
        while self.par[x] != x:
            x = self.par[x]
        while self.par[a] != x:
            self.par[a], a = x, self.par[a]
        return x

    def union(self, x, y):
        xr, yr = self.find(x), self.find(y)
        if xr == yr:
            return False
        else:
            if self.size[xr] < self.size[yr]:
                xr, yr = yr, xr
            self.par[yr] = xr
            self.size[xr] += self.size[yr]
            self.size[yr] = 0
            return True

f = open("day18.input")

arr = []

d = defaultdict(lambda: defaultdict(lambda: defaultdict(lambda: False)))

a,b,c = -math.inf, -math.inf, -math.inf

for line in f:
  line = line.strip()
  sp = map(int, line.split(","))
  x,y,z = sp

  a = max(a, x)
  b = max(b, y)
  c = max(c, z)

  arr.append( [x,y,z] )
  d[x][y][z] = True

f.close()

DIRS = [ (0,0,1), (0,1,0), (1,0,0), (0,0,-1), (0,-1,0), (-1, 0, 0) ]
res = 0
cands = set()

for x, y, z in arr:
  for dx, dy, dz in DIRS:
    nx, ny, nz = x+dx, y+dy, z+dz
    if not d[nx][ny][nz]:
      cands.add( (nx, ny, nz) )
      res += 1

MAX_N = max(a,b,c) + 1
dsu = DSU_Iterative_By_Size(MAX_N*MAX_N*MAX_N)

for i in range(MAX_N): # all start from 0
  for j in range(MAX_N):
    for k in range(MAX_N):
      if d[i][j][k]: continue
      # union every two adjacent empty
      idx = i*(MAX_N*MAX_N) + j*MAX_N + k
      for dx, dy, dz in DIRS:
        nx, ny, nz = i+dx, j+dy, k+dz
        if not ( 0<=nx< MAX_N and 0<=ny<MAX_N and 0<=nz<MAX_N ):
          continue
        if d[nx][ny][nz]: continue
        nidx = nx*(MAX_N*MAX_N) + ny*MAX_N + nz
        dsu.union(idx, nidx)

# note: could be part of a larger bubble.

def is_free(i, j, k):
  idx = i*(MAX_N*MAX_N) + j*MAX_N + k
  return dsu.find(idx) == dsu.find(0)

trapped = defaultdict(lambda: defaultdict(lambda: defaultdict(lambda: False)))

for x, y, z in list(cands):
  if not is_free(x, y, z):
    trapped[x][y][z] = True

res = 0
for x, y, z in arr:
  for dx, dy, dz in DIRS:
    nx, ny, nz = x+dx, y+dy, z+dz
    if not d[nx][ny][nz] and not trapped[nx][ny][nz]:
      cands.add( (nx, ny, nz) )
      res += 1

print(res)