aoc2024day20.py

# === Parsing ===

valid = set()
start = None
end = None

with open("input.txt", "r") as f:
    for y, line in enumerate(f):
        for x, char in enumerate(line.strip()):
            if char in ".SE":
                valid.add((x, y))
                if char == "S":
                    start = (x, y)
                elif char == "E":
                    end = (x, y)

assert start is not None
assert end is not None

print(start, end)


# === Part 1 ===
DIRECTIONS = [
    (0, 1),
    (1, 0),
    (0, -1),
    (-1, 0),
]


def find_legal_path(start, end, valid):
    path = [start]
    while path[-1] != end:
        px, py = path[-1]
        for dx, dy in DIRECTIONS:
            nx, ny = px + dx, py + dy
            if (nx, ny) in valid and (len(path) == 1 or path[-2] != (nx, ny)):
                path.append((nx, ny))
                break

    return path


def find_cheats(path, save_at_least):
    path_with_costs = list(enumerate(path))
    pos_costs = {pos: cost for cost, pos in path_with_costs}
    cheats_that_worked = set()

    for cost, pos in path_with_costs:
        px, py = pos
        for dx, dy in DIRECTIONS:
            nx, ny = px + 2 * dx, py + 2 * dy
            if (nx, ny) in pos_costs:
                cost_saved = pos_costs[nx, ny] - cost - 2
                if cost_saved >= save_at_least:
                    cheats_that_worked.add((pos, (nx, ny)))

    return cheats_that_worked


path = find_legal_path(start, end, valid)
print(len(find_cheats(path, 100)))
print("---")


# === Part 2 ===
def find_cheats_far(path, save_at_least, max_cheat_dist):
    cheats_that_worked = set()

    for from_idx, from_ in enumerate(path):
        offset_idx = from_idx + save_at_least
        for to_idx, to_ in enumerate(path[offset_idx:], start=offset_idx):
            fx, fy = from_
            tx, ty = to_
            dist = abs(fx - tx) + abs(fy - ty)
            if dist > max_cheat_dist:
                continue
            cost_saved = to_idx - from_idx - dist
            if cost_saved >= save_at_least:
                cheats_that_worked.add((from_, to_))

    return cheats_that_worked


print(len(find_cheats_far(path, 100, 2)))
print(len(find_cheats_far(path, 100, 20)))
print("---")


# === Part 2, alternative ===
def find_cheats_far2(path, save_at_least, max_cheat_dist):
    pos_costs = {pos: cost for cost, pos in enumerate(path)}
    cheats_that_worked = set()

    for cost, (px, py) in enumerate(path):
        for dx in range(-max_cheat_dist, max_cheat_dist + 1):
            for dy in range(-(max_cheat_dist - abs(dx)), max_cheat_dist - abs(dx) + 1):
                nx, ny = px + dx, py + dy
                if (nx, ny) in pos_costs:
                    dist = abs(dx) + abs(dy)
                    cost_saved = pos_costs[nx, ny] - cost - dist
                    if cost_saved >= save_at_least:
                        cheats_that_worked.add(((px, py), (nx, ny)))

    return cheats_that_worked


print(len(find_cheats_far2(path, 100, 2)))
print(len(find_cheats_far2(path, 100, 20)))