day20.um

type Pos = struct {
    x, y: int
}

fn posmove(pos: Pos): [4]Pos {
    return {
        {pos.x + 1, pos.y},
        {pos.x - 1, pos.y},
        {pos.x, pos.y + 1},
        {pos.x, pos.y - 1}
    }
}

fn istaken(m: []str, pos: Pos): bool {
    if pos.x < 0 || pos.x >= len(m[0]) || pos.y < 0 || pos.y >= len(m) {
        return true
    }

    return m[pos.y][pos.x] == '#'
}

type Step = struct {
    pos: Pos
    cost: int
}

fn findDefaultPath(m: []str, start: Pos, end: Pos): (map[Pos]Step, int) {
    steps := 0
    trace := map[Pos]Step{}

    for !(start.x == end.x && start.y == end.y) {
        moves := posmove(start)
        for _, newPos in moves {
            if validkey(trace, newPos) {
                continue
            }

            if !istaken(m, newPos) {
                steps++
                trace[start] = {newPos, steps}
                start = newPos
                break
            }
        }
    }

    trace[start] = {start, steps+1}

    return trace, steps
}

fn findAllShortcuts(m: []str, trace: map[Pos]Step, start: Pos, lim: int, radius: int = 2): map[int]int {
    scores := map[int]int{}

    for i := 0; i < lim; i++ {
        for j := -radius; j <= radius; j++ {
            w := radius-abs(j)
            for k := 0; k <= w*2; k++ {
                x, y := w-k, j

                dist := abs(x)+abs(y)
                newPos := Pos{start.x + x, start.y + y}
                if !istaken(m, newPos) {
                    c := lim-(trace[newPos].cost-trace[start].cost)+dist
                    if c < lim {
                        scores[c]++
                    }
                }
            }
        }

        start = trace[start].pos
    }

    return scores
}

fn main() {
    m := []str{}
    for l := ""; scanf("%s", &l) == 1 {
        m = append(m, l)
    }

    start, end := Pos{0, 0}, Pos{0, 0}

    for i, l in m {
        for j, c in l {
            if c == 'S' {
                start = Pos{j, i}
            }
            if c == 'E' {
                end = Pos{j, i}
            }
        }
    }

    defaultPath, lim := findDefaultPath(m, start, end)
    top := 100
    if len(m[0]) < 100 { // example
        top = 50
    }

    p1, p2 := 0, 0

    tally1 := findAllShortcuts(m, defaultPath, start, lim)
    for k, v in tally1 {
        if lim-k >= top {
            p1 += v
        }
    }
    printf("Part 1: %v\n", p1)

    tally2 := findAllShortcuts(m, defaultPath, start, lim, 20)
    for k, v in tally2 {
        if lim-k >= top {
            p2 += v
        }
    }
    printf("Part 2: %v\n", p2)
}