day9.py

FILE = 'input.txt'

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

def parse_moves():
    moves = []
    with open(FILE, 'r') as f:
        for line in f.readlines():
            direction, num = line.strip().split(' ')
            moves.append((direction, int(num)))

    return moves

def make_move(pos, direction):
    delta = DIRECTION_TO_DELTA[direction]
    return (pos[0] + delta[0], pos[1] + delta[1])

def touching(head_pos, tail_pos) -> bool:
    x_delta = head_pos[0] - tail_pos[0]
    y_delta = head_pos[1] - tail_pos[1]

    return abs(x_delta) < 2 and abs(y_delta) < 2
    
def move_tail(head_pos, tail_pos):
    if head_pos == tail_pos:
        return tail_pos

    x_delta = head_pos[0] - tail_pos[0]
    y_delta = head_pos[1] - tail_pos[1]

    new_tail_x = tail_pos[0]
    new_tail_y = tail_pos[1]

    # If in same row, check if H got 2 ahead/behind of T
    if head_pos[1] == tail_pos[1]:
        if x_delta == 2:
            new_tail_x += 1
        elif x_delta == -2:
            new_tail_x -= 1
    elif head_pos[0] == tail_pos[0]:  # same col, same check
        if y_delta == 2:
            new_tail_y += 1
        elif y_delta == -2:
            new_tail_y -= 1
    elif not touching(head_pos, tail_pos):  #  move tail towards head diagonally 1 step
        new_tail_x += x_delta // abs(x_delta)
        new_tail_y += y_delta // abs(y_delta)

    return (new_tail_x, new_tail_y)


def part1():
    head_pos = (0, 0)
    tail_pos = (0, 0)
    visited = set([tail_pos])  # will have (x, y) tuple elements

    for move in parse_moves():
        direction, num = move
        for _ in range(num):
            head_pos = make_move(head_pos, direction)
            tail_pos = move_tail(head_pos, tail_pos)
            visited.add(tail_pos)

    return len(visited)

NUM_KNOTS = 10

def part2():
    knots = [(0, 0) for _ in range(NUM_KNOTS)]
    visited = set([knots[-1]])

    for move in parse_moves():
        direction, num = move
        for _ in range(num):
            knots[0] = make_move(knots[0], direction)
            for i in range(1, NUM_KNOTS):
                knots[i] = move_tail(knots[i-1], knots[i])

            visited.add(knots[-1])

    return len(visited)