adventofcode2022/day09/day09.py

def go(p, dir):
    x,y = p
    if   dir == "U": y += 1
    elif dir == "D": y -= 1
    elif dir == "L": x -= 1
    elif dir == "R": x += 1
    return x,y

def follow(h, t):
    if abs(h[0] - t[0]) > 1 or abs(h[1] - t[1]) > 1:
        x,y = t
        x += sgn(h[0] - t[0])
        y += sgn(h[1] - t[1])
        t = x,y
    return t

def sgn(x):
    return (x > 0) - (x < 0)

def solve(n):
    R = [(0,0)]*n
    visited = set([R[-1]])
    for line in open("input"):
        dir, count = line.split()
        for _ in range(int(count)):
            R[0] = go(R[0], dir)
            for i in range(1, len(R)):
                t = follow(R[i-1], R[i])
                if t == R[i]:
                    # optimization: if a knot doesn't move,
                    # none of the knots after it will either
                    break
                R[i] = t
            visited.add(R[-1])
            #print(R)
    return len(visited)

A1 = solve(2)
print(A1)

A2 = solve(10)
print(A2)