AoC 2024 day 20

P=puzzle.splitlines()
D=[-1j,1,1j,-1]                                                 # move directions
F=set()                                                         # set of free positions
for j,r in enumerate(P):                                        # for each row
    for i,c in enumerate(r):                                    # for each column
        p=i+1j*j                                                # position in the maze
        if c=='S': S=p                                          # start position
        if c=='E': E=p                                          # end position
        if c!='#': F.add(p)                                     # memorize free positions

def maze():                                                     # find the shortest path from start to end
    q=[(0,[S])]; bc=1e9; v=set()                                # queue the start cost/path, best cost, visited positions
    while(q):                                                   # while unexplored positions remain
        c,pp=q.pop(0);p=pp[-1]                                  # unqueue the 1st cost/path for exploration, last position
        if c<bc:                                                # no exploration needed if cost is already higher than best
            if p==E: return(pp)                                 # if we reached the end, return the path solution
            for np in [p+d for d in D]:                         # else check the neighbour positions
                if np in F and np not in v:                     # for the free ones not already visited
                    v.add(np);q.append((c+1,pp+[np]))           # mark them as visited and queue them for later exploration

def cheat(cd,sv):                                               # how many ways to save time (sv) if collision disabled steps (cd) are allowed
    pl=maze(); ch=0                                             # resolve the path list of the maze
    pi=dict({p:i for i,p in enumerate(pl)})                     # position index for each path step
    for p,n in pi.items():                                      # for each path position from start
        for dj in range(-cd,cd+1):                              # for each next position within a Manhattan distance equal to cd
            adj=abs(dj)                                         # such positions are in a diamond shape around the current position
            for di in range(adj-cd,cd+1-adj):                   # typically, if  -d<= dy <= d , then  |y|-d <= dx <= d-|y| 
                adi=abs(di); np=p+di+1j*dj; md=adi+adj          # next position and its Manhattan distance from current position
                if md>1 and np in pi:                           # need to move at least 2 steps to cross a wall, else there would be no cheating
                    if pi[np]-n-md>=sv: ch+=1                   # count the nbr of cheating path that saves us enough time, from this position
    return(ch)                                                  # display the total number

def part1():
    print(cheat(2,1 if _F==1 else 100))                         # nbr of cheats saving min 100 picosec, if collision is disabled during max 2 steps

def part2():
    print(cheat(20,50 if _F==1 else 100))                       # nbr of cheats saving min 100 picosec, if collision is disabled during max 20 steps