Matthew Gretton-Dann
4bc7815851
We filter out already visited states when adding them to the work list. We also don't sort the work list as we're already generating it in a sorted (by cost) order.
97 lines
3.0 KiB
OCaml
97 lines
3.0 KiB
OCaml
let find grid c =
|
|
match Aoc.Grid.idx_from_opt grid 0 c with
|
|
| None -> failwith "find"
|
|
| Some idx -> Aoc.Grid.pos_of_idx grid idx
|
|
|
|
let input_of_file fname =
|
|
let grid = Aoc.Grid.of_file fname in
|
|
let start_pos = find grid 'S' in
|
|
(grid, start_pos)
|
|
|
|
let rec dijkstra visit check_end states =
|
|
let compare_costs (lhs, _) (rhs, _) = compare lhs rhs in
|
|
match states with
|
|
| [] -> None
|
|
| (cost, state) :: t ->
|
|
if check_end state then Some (cost, state, t)
|
|
else
|
|
let new_states = visit cost state |> List.merge compare_costs t in
|
|
dijkstra visit check_end new_states
|
|
|
|
let visited_idx grid ((dx, dy), p) =
|
|
let add =
|
|
match (dx, dy) with
|
|
| 1, 0 -> 0
|
|
| 0, 1 -> 1
|
|
| -1, 0 -> 2
|
|
| 0, -1 -> 3
|
|
| _ -> failwith "visited_idx"
|
|
in
|
|
(Aoc.Grid.idx_of_pos grid p * 4) + add
|
|
|
|
let visit grid visited_grid cost state =
|
|
let (dx, dy), ((x, y) as p) = List.hd state in
|
|
let has_visited = visited_grid.(visited_idx grid (List.hd state)) in
|
|
if has_visited then []
|
|
else if Aoc.Grid.get_by_pos grid p = '#' then []
|
|
else (
|
|
visited_grid.(visited_idx grid (List.hd state)) <- true;
|
|
[
|
|
(cost + 1, ((dx, dy), (x + dx, y + dy)) :: state);
|
|
(cost + 1000, ((-dy, dx), p) :: state);
|
|
(cost + 1000, ((dy, -dx), p) :: state);
|
|
])
|
|
|
|
let has_visited grid visited_grid (cost, state) =
|
|
visited_grid.(visited_idx grid (List.hd state)) < cost
|
|
|
|
let visit_max grid visited_grid cost state =
|
|
let (dx, dy), ((x, y) as p) = List.hd state in
|
|
if has_visited grid visited_grid (cost, state) then []
|
|
else if Aoc.Grid.get_by_pos grid p = '#' then []
|
|
else (
|
|
Printf.printf "%d (%d, %d)\n" cost x y;
|
|
flush stdout;
|
|
visited_grid.(visited_idx grid (List.hd state)) <- cost;
|
|
[
|
|
(cost + 1, ((dx, dy), (x + dx, y + dy)) :: state);
|
|
(cost + 1000, ((-dy, dx), p) :: state);
|
|
(cost + 1000, ((dy, -dx), p) :: state);
|
|
]
|
|
|> List.filter (fun x -> not (has_visited grid visited_grid x)))
|
|
|
|
let[@warning "-32"] print_point (x, y) = Printf.printf "(%d, %d)" x y
|
|
|
|
let check_end grid state =
|
|
let _, p = List.hd state in
|
|
Aoc.Grid.get_by_pos grid p = 'E'
|
|
|
|
let check_end2 grid _ state = check_end grid state
|
|
|
|
let part1 (grid, start_pos) =
|
|
let visited_grid = Array.make (Aoc.Grid.length grid * 4) false in
|
|
match
|
|
dijkstra (visit grid visited_grid) (check_end grid)
|
|
[ (0, [ ((1, 0), start_pos) ]) ]
|
|
with
|
|
| None -> failwith "part"
|
|
| Some (cost, _, _) -> cost
|
|
|
|
let part2 (grid, start_pos) =
|
|
let cost = part1 (grid, start_pos) in
|
|
let visited_grid = Array.make (Aoc.Grid.length grid * 4) cost in
|
|
let rec impl acc lst =
|
|
match
|
|
dijkstra (visit_max grid visited_grid) (check_end2 grid visited_grid) lst
|
|
with
|
|
| None -> acc
|
|
| Some (_, states, remainder) ->
|
|
List.iter (fun x -> visited_grid.(visited_idx grid x) <- 0) states;
|
|
impl (states :: acc) remainder
|
|
in
|
|
impl [] [ (0, [ ((1, 0), start_pos) ]) ]
|
|
|> List.concat |> List.map snd |> List.sort_uniq compare |> List.length
|
|
|
|
let _ =
|
|
Aoc.main input_of_file [ (string_of_int, part1); (string_of_int, part2) ]
|