Advent of code 2020, Day 24, Python3.9

day24.py

from sys import stdin
from functools import reduce
from collections import Counter

DIRS = ['e', 'se', 'sw', 'w', 'nw', 'ne']


def parse(l):
    res = []
    while l:
        for d in DIRS:
            if l.startswith(d):
                res += [d]
                l = l[len(d):]
    return res


def jump(xyz, dirs):
    def step(x, y, z, d): return {
        'e':  (x-1, y+1, z),
        'w':  (x+1, y-1, z),
        'se': (x, y+1, z-1),
        'sw': (x+1, y, z-1),
        'nw': (x, y-1, z+1),
        'ne': (x-1, y, z+1)
    }[d]
    return reduce(lambda xyz, d: step(*xyz, d),
                  dirs, xyz)


def flip(blacks):
    neighbours = Counter(jump(t, [d]) for d in DIRS for t in blacks)
    blacks = set(t for t, val in neighbours.items()
                 if val == 2 or t in blacks and val == 1)
    return blacks


blacks = set()
for d in [parse(l.strip()) for l in stdin]:
    if (t := jump((0, 0, 0), d)) in blacks:
        blacks.remove(t)
    else:
        blacks.add(t)

# Part 1
print(len(blacks))

# Part 2
for _ in range(100):
    neighbours = Counter(jump(t, [d]) for d in DIRS for t in blacks)
    blacks = set(t for t, val in neighbours.items()
                 if val == 2 or t in blacks and val == 1)
print(len(blacks))