aoc/2022/9/solution.py

#!/usr/bin/python3
import sys
from dataclasses import dataclass

@dataclass
class Pos:
    x: int
    y: int

    def move(self, pos):
        self.x += pos.x
        self.y += pos.y
    
    def behind(self, pos):
        if abs(self.x - pos.x) <= 1 and abs(self.y - pos.y) <= 1:
            return
        elif abs(self.x - pos.x) > 1 and abs(self.y - pos.y) > 1:
            x_sign = 1 if self.x - pos.x > 1 else -1
            y_sign = 1 if self.y - pos.y > 1 else -1
            self.x = pos.x + x_sign
            self.y = pos.y + y_sign
        elif abs(self.x - pos.x) > 1:
            sign = 1 if self.x - pos.x > 1 else -1
            self.y = pos.y
            self.x = pos.x + sign
        elif abs(self.y - pos.y) > 1:
            sign = 1 if self.y - pos.y > 1 else -1
            self.x = pos.x
            self.y = pos.y + sign
    
    def as_struct(self):
        return (self.x, self.y)

# part 1
directions = { 'U': Pos(0, 1), 'D': Pos(0, -1), 'L': Pos(-1, 0), 'R': Pos(1, 0) }

def visualize(visited, max_width=180):
    y_s = [y for x, y in visited]
    y_start, y_end = min(y_s), max(y_s)
    x_s = [x for x, y in visited]
    x_start, x_end = min(x_s), max(x_s)
    for y in range(y_start, y_end + 1)[::-1]:
        for x in range(x_start, min(x_start + max_width, x_end) + 1):
            print('#' if (x, y) in visited else '.', end='')
        print()

def count_positions(visited):
    return len(visited)

def solution(n):
    head = Pos(0, 0)
    tails = [Pos(0, 0) for i in range(n)]
    visited = set()
    for line in sys.stdin:
        s = line.strip().split(' ')
        direction, count = s[0], int(s[1])
        for i in range(count):
            head.move(directions[direction])
            tails[0].behind(head)
            for i in range(1, n):
                tails[i].behind(tails[i - 1])
            if tails[-1].as_struct() not in visited:
                visited.add(tails[-1].as_struct())
            #print(head, tails)
    #print(sorted(visited))
    visualize(visited)
    print(f'{count_positions(visited)} visited')

if sys.argv[1] in '1':
    solution(1)
else:
    solution(9)