73 lines
2.6 KiB
OCaml
73 lines
2.6 KiB
OCaml
module CharMap = Map.Make (Char)
|
|
|
|
(** [get_station_indices map] returns a list of pairs mapping station ID to the
|
|
indices in [map] where there is a station with that ID. *)
|
|
let get_station_indices map =
|
|
let rec impl acc idx =
|
|
if idx >= Aoc.Grid.length map then acc
|
|
else if Aoc.Grid.get_by_idx map idx = '.' then impl acc (idx + 1)
|
|
else
|
|
let station = Aoc.Grid.get_by_idx map idx in
|
|
let update_fn lst =
|
|
match lst with None -> Some [ idx ] | Some t -> Some (idx :: t)
|
|
in
|
|
impl (CharMap.update station update_fn acc) (idx + 1)
|
|
in
|
|
impl CharMap.empty 0 |> CharMap.to_list
|
|
|
|
(** Generate antinodes for part 1. *)
|
|
let get_antinodes1 acc _ (px, py) (px', py') =
|
|
let dx = px' - px in
|
|
let dy = py' - py in
|
|
(px - dx, py - dy) :: (px' + dx, py' + dy) :: acc
|
|
|
|
(** [add_antinodes lst map pos vel] adds antinodes at [pos + n * vel] to [lst]
|
|
for all non-negative [n] that are valid positions in [map]. *)
|
|
let rec add_antinodes lst map (x, y) (dx, dy) =
|
|
if Aoc.Grid.pos_is_valid map (x, y) then
|
|
add_antinodes ((x, y) :: lst) map (x + dx, y + dy) (dx, dy)
|
|
else lst
|
|
|
|
(** Generate antinodes for part 2. *)
|
|
let get_antinodes2 acc map (px, py) (px', py') =
|
|
let dx = px' - px in
|
|
let dy = py' - py in
|
|
let acc' = add_antinodes acc map (px, py) (dx, dy) in
|
|
let acc'' = add_antinodes acc' map (px, py) (-dx, -dy) in
|
|
acc''
|
|
|
|
(** [process_stations map fn stations] generates a list of all antinodes for the
|
|
stations in [stations] on the map [map]. [fn acc map p p'] is called to
|
|
generate the antinode list for each pair of stations [p] and [p']. It should
|
|
add the positions of antinodes to the list [acc]. *)
|
|
let process_stations map fn stations =
|
|
let rec impl2 acc p t =
|
|
match t with
|
|
| [] -> acc
|
|
| h :: t -> impl2 (fn acc map p (Aoc.Grid.pos_of_idx map h)) p t
|
|
in
|
|
let rec impl acc = function
|
|
| [] -> acc
|
|
| h :: t -> impl (impl2 acc (Aoc.Grid.pos_of_idx map h) t) t
|
|
in
|
|
List.map (impl []) stations
|
|
|
|
(** [part antifn map station_indices] process all the stations in
|
|
station_indices calling [antifn acc map p p'] on all stations. Here [acc] is
|
|
a list of antinodes which [antifn] should update and return, [p] and [p']
|
|
are positions of stations to generate antinodes for. *)
|
|
let part antifn map =
|
|
get_station_indices map
|
|
|> List.map snd (* we do not care about the station IDs *)
|
|
|> process_stations map antifn
|
|
|> List.concat
|
|
|> List.filter (Aoc.Grid.pos_is_valid map)
|
|
|> List.sort_uniq Stdlib.compare
|
|
|> List.length
|
|
|
|
let _ =
|
|
Aoc.main Aoc.Grid.of_file
|
|
[
|
|
(string_of_int, part get_antinodes1); (string_of_int, part get_antinodes2);
|
|
]
|