aoc 2023 day 22 pt 2.py

aoc 2023 day 22 pt 2.py

from util import ints
from collections import defaultdict
from functools import cmp_to_key
from copy import deepcopy, copy

prob = "22"
f = open(prob, "r")

lines = []

for line in f.readlines():
    line = line.strip()
    lines.append(line)
   
id = 0
pts = defaultdict(set)
for line in lines:
    p1, p2 = line.split("~")
    x1,y1,z1 = ints(p1)
    x2,y2,z2 = ints(p2)

    dx = x2-x1
    dy = y2-y1
    dz = z2-z1

    vec = (dx,dy,dz)
    t = max([abs(dx),abs(dy),abs(dz)])
    vec = (int(vec[0]/max(1,abs(vec[0]))), int(vec[1]/max(1,abs(vec[1]))), int(vec[2]/max(1,abs(vec[2]))))

    pts[id].add((x1,y1,z1))
    while t > 0:
        p = (x1,y1,z1)
        while p != (x2,y2,z2):
            p = (p[0]+vec[0],p[1]+vec[1],p[2]+vec[2])
            pts[id].add(p)
        pts[id].add(p)
        t -= 1
    id += 1

while True:
    moved = False
    for id in pts:
        newpts = [(p[0],p[1],p[2]-1) for p in pts[id]]
        if all([pt[2] >= 1 for pt in newpts]):
            newpts = set(newpts)
            canmove = True
            for p in newpts:
                if not canmove:
                    break
                for id2 in pts:
                    if id2 != id and p in pts[id2]:
                        canmove = False
                        break

            if canmove:
                pts[id] = newpts
                moved = True

    if not moved:
        break

supports = defaultdict(set)

m = {}
for id in pts:
    for pt in pts[id]:
        m[pt] = id

indegree = defaultdict(set)
for id in pts:
    for pt in pts[id]:
        x,y,z = pt
        up = (x,y,z+1)
        if up in m and m[up] != id: 
            supports[id].add(m[up])
            indegree[m[up]].add(id)

indegree_real = { k: len(indegree[k]) for k in indegree }
indegree = indegree_real

ans1 = [1 for i in pts if all([indegree[k] > 1 for k in supports[i]])]
print("Pt1", sum(ans1))

ans2=0
for sid in pts:
    indegree_cpy = copy(indegree)
    collapsed = 0
    Q = [sid]
    while len(Q) > 0:
        v = Q[0]
        del Q[0]
        for n in supports[v]:
            indegree_cpy[n] -= 1
            if indegree_cpy[n] == 0:
                collapsed += 1
                Q.append(n)
    ans2 += collapsed

print("Pt2", ans2)