aoc/2017/22_virus.py

#!/usr/bin/env python3
import sys


class PartOne:
    def __init__(self, grid):
        self.infected = set()
        for y, line in enumerate(grid.split('\n')):
            for x, cell in enumerate(line):
                if cell == '#':
                    self.infected.add((y, x))
        self.y = y // 2
        self.x = x // 2
        self.dy = -1
        self.dx = 0

    def turn_left(self):
        self.dy, self.dx = -self.dx, self.dy

    def turn_right(self):
        self.dy, self.dx = self.dx, -self.dy

    def advance(self):
        pos = self.y, self.x
        if pos in self.infected:
            self.infected.remove(pos)
            self.turn_right()
            infection = False
        else:
            self.infected.add(pos)
            self.turn_left()
            infection = True

        self.y += self.dy
        self.x += self.dx
        return infection


CLEAN, FLAGGED, INFECTED, WEAKENED = range(4)


class PartTwo(PartOne):
    def __init__(self, grid):
        self.nodes = {}
        for y, line in enumerate(grid.split('\n')):
            for x, cell in enumerate(line):
                if cell == '#':
                    self.nodes[(y, x)] = INFECTED
        self.y = y // 2
        self.x = x // 2
        self.dy = -1
        self.dx = 0

    def at(self, y, x):
        return self.nodes.get((y, x), CLEAN)

    def turn_back(self):
        self.dy, self.dx = -self.dy, -self.dx

    def advance(self):
        state = self.at(self.y, self.x)

        if state == CLEAN:
            self.turn_left()
        elif state == INFECTED:
            self.turn_right()
        elif state == FLAGGED:
            self.turn_back()

        self.nodes[(self.y, self.x)] = (self.at(self.y, self.x) + 3) % 4

        self.y += self.dy
        self.x += self.dx
        return state == WEAKENED


grid = sys.stdin.read().rstrip()

state = PartOne(grid)
print(sum(int(state.advance()) for i in range(10000)))

state = PartTwo(grid)
print(sum(int(state.advance()) for i in range(10000000)))