aoc_2021_15

dijkstra.ml

module Coordinate = struct
  type t = { x : int; y : int; value : int }

  let make ~x ~y value = { x; y; value }

  let compare t1 t2 =
    match compare t1.x t2.x with 0 -> compare t1.y t2.y | _ as x -> x
end

module PQ = Psq.Make (Coordinate) (Int)
module CMap = CCMap.Make (Coordinate)

let neighbour ~x ~y matrix =
  let w = Array.length matrix.(0) in
  let h = Array.length matrix in
  let bounds (x, y) = x >= 0 && y >= 0 && x < h && y < w in
  [ (x, succ y); (succ x, y); (x, pred y); (pred x, y) ]
  |> List.filter ~f:bounds

let dijkstra ~source ~target matrix =
  let pq = ref (PQ.add source 0 PQ.empty) in
  let distance = ref (CMap.add source 0 CMap.empty) in

  for y = 0 to pred @@ Array.length matrix do
    for x = 0 to pred @@ Array.length matrix.(0) do
      let c = Coordinate.make ~x ~y matrix.(y).(x) in
      distance := CMap.add c Int.max_int !distance
    done
  done;

  while not (PQ.is_empty !pq) do
    let (u, u_dist), npq =
      match PQ.pop !pq with
      | None -> failwith "this loop only runs when pq isn't empty..."
      | Some (c, npq) -> (c, npq)
    in
    let () = pq := npq in
    let neighbours = Coordinate.(neighbour ~x:u.x ~y:u.y matrix) in

    let check_neighbour (x, y) =
      let v = Coordinate.make ~x ~y matrix.(y).(x) in
      let v_dist = Coordinate.(v.value) in

      match CMap.get v !distance with
      | None -> ()
      | Some old_distance ->
          let new_distance = u_dist + v_dist in
          if old_distance > new_distance then
            let () = distance := CMap.add v new_distance !distance in
            let () = pq := PQ.add v new_distance !pq in
            ()
          else ()
    in
    let () = List.iter ~f:check_neighbour neighbours in
    ()
  done;
  CMap.get target !distance