aoc-2022/9.py

lines = [(line[0], int(line[2:])) for line in open("9.input")]

DIR = {'R': (0, 1), 'L': (0, -1), 'U': (-1, 0), 'D': (1, 0)}

#SNEK_LEN = 2  # Part 1
SNEK_LEN = 10  # Part 2
snek = [(0, 0) for _ in range(0, SNEK_LEN)]

def add(a, b):
    return (a[0] + b[0], a[1] + b[1])

def show_snek(snek):
    for x in range(-20, 20):
        for y in range(-20, 20):
            if (x, y) in snek:
                print("#", end='')
            elif (x, y) == (0, 0):
                print("s", end='')
            else:
                print(".", end='')
        print("")

visited = set()
visited.add(snek[-1])

for direction, steps in lines:
    for _ in range(0, steps):
        # Move head
        snek[0] = add(snek[0], DIR[direction])

        # Follow with tail
        for head_idx in range(0, len(snek) - 1):
            head = snek[head_idx]
            tail = snek[head_idx+1]

            follow = [0, 0]
            follow[0] = head[0] - tail[0]
            follow[1] = head[1] - tail[1]
            dist = sum(map(abs, follow))
            if dist >= 3:
                follow[0] = follow[0] // abs(follow[0])
                follow[1] = follow[1] // abs(follow[1])
            elif abs(follow[0]) == 2:
                follow[0] = follow[0] // 2
            elif abs(follow[1]) == 2:
                follow[1] = follow[1] // 2
            else:
                follow = [0, 0]

            snek[head_idx+1] = add(snek[head_idx+1], follow)

        visited.add(snek[-1])

print(len(visited))