Untitled

from advent_of_code.lib.all import *


def get_input(file):
    raw = aoc.read_input(2024, 20, file)
    grid = aoc_parse.as_lines(raw)
    m = {r+1j*c:x for r, line in enumerate(grid) for c, x in enumerate(line)}
    return m


def shortest(m, start, end):
    q = [(start, 0)]
    seen = set([start])
    while q:
        p, l = q.pop()
        if p == end:
            return l
        # visit neighbours
        for step in (1, -1, 1j, -1j):
            new = p + step
            if new in m and m[new] != "#" and new not in seen:
                q.append((new, l+1))
                seen.add(new)
    return None


def bfs(m, start):
    res = {start:0}
    q = [(0, t:=0, start)]
    seen = set([start])
    while q:
        l, _, p = heappop(q)
        res[p] = l
        # visit neighbours
        for step in (1, -1, 1j, -1j):
            new = p + step
            if new in m and m[new] != "#" and new not in seen:
                heappush(q, (l+1, t:=t+1, new))
                seen.add(new)
    return res


def dfs(m, start, end_dist, target, cheat_length):
    dirs = (1, -1, 1j, -1j)
    q = [(start, 0)]
    seen = {start}
    # set of cheats as (start pos, end pos)
    cheats = set()
    while q:
        # regular dfs
        p, l = q.pop()
        for step in dirs:
            new = p + step
            if m[new] != "#" and l+1 <= target and new not in seen:
                q.append((new, l+1))
                seen.add(new)
        # use cheat: dfs inside dfs
        q2 = [(p, 0)]
        seen2 = {p}
        while q2:
            p2, l2 = q2.pop()
            if m[p2] != "#" and end_dist[p2] + l + l2 <= target and (p, p2) not in cheats:
                cheats.add((p, p2))
            for step2 in dirs:
                new2 = p2 + step2
                if new2 in m and new2 not in seen2 and l2 +1 <= cheat_length:
                    q2.append((new2, l2+1))
                    seen2.add(new2)
    return len(cheats)


@aoc.pretty_solution(1)
def part1(m):
    start, end = map(lambda x: next(p for p in m if m[p] == x), ["S", "E"])
    target = shortest(m, start, end) - 100
    end_dist = bfs(m, end)
    return dfs(m, start, end_dist, target, 2)


@aoc.pretty_solution(2)
def part2(m):
    start, end = map(lambda x: next(p for p in m if m[p] == x), ["S", "E"])
    target = shortest(m, start, end) - 50
    end_dist = bfs(m, end)
    return dfs(m, start, end_dist, target, 20)


def test():
    data = get_input("input.txt")
    part1(deepcopy(data))
    part2(data)
    # assert part1(deepcopy(data)) == 1369
    # assert part2(data) == 
    print("Test OK")


if __name__ == "__main__":
    test()