57 lines
2.2 KiB
OCaml
57 lines
2.2 KiB
OCaml
module IntMap = Map.Make (Int)
|
|
|
|
(** [apply_n n fn arg] is equivalent to [(fn (fn ... (fn (fn arg))))] where [fn]
|
|
is called [n] times.*)
|
|
let rec apply_n n fn arg = if n <= 0 then arg else apply_n (n - 1) fn (fn arg)
|
|
|
|
(** [update_count n o] returns [Some n] if [o] is [None], or [Some (n + x)] if
|
|
[o] is [Some x]. *)
|
|
let update_count n = function None -> Some n | Some x -> Some (x + n)
|
|
|
|
(** [map_of_ints lst] returns an [IntMap] with key-value pairs counting how many
|
|
times each integer appears in [lst]. *)
|
|
let map_of_ints =
|
|
let rec impl acc = function
|
|
| [] -> acc
|
|
| h :: t -> impl (IntMap.update h (update_count 1) acc) t
|
|
in
|
|
impl IntMap.empty
|
|
|
|
(** [map_stone id n map] calculates how a collection of stones changes in a
|
|
blink, updating [map] with the result. [id] is the ID of the stone to
|
|
update, [n] is how many stones there are with that ID. *)
|
|
let map_stone id n acc =
|
|
match id with
|
|
| 0 -> IntMap.update 1 (update_count n) acc
|
|
| x when Aoc.digits10 x mod 2 = 0 ->
|
|
let pow = Aoc.pow10 (Aoc.digits10 x / 2) in
|
|
let acc = IntMap.update (x / pow) (update_count n) acc in
|
|
let acc = IntMap.update (x mod pow) (update_count n) acc in
|
|
acc
|
|
| x -> IntMap.update (x * 2024) (update_count n) acc
|
|
|
|
(** [calc_blink map] calculates how a collection of stones changes in a blink,
|
|
returning the result.
|
|
|
|
This improves performance because we find that the transformation stones go
|
|
through ends up producing repeated numbers. e.g.: 0 -> 1 -> 2024 -> 20 24 ->
|
|
2 0 2 4 which has two 2s in it.
|
|
|
|
We also note that despite the problem description saying stones stay in
|
|
order, the result we are asked for (number of stones) does not require them
|
|
to be in order. *)
|
|
let calc_blink map = IntMap.fold map_stone map IntMap.empty
|
|
|
|
(** [part n lst] returns the number of stones after [n] blinks, given an initial
|
|
list, [lst], of stone IDs. *)
|
|
let part n lst =
|
|
let map = map_of_ints lst |> apply_n n calc_blink in
|
|
IntMap.fold (fun _ v acc -> v + acc) map 0
|
|
|
|
(** [ints_of_file fname] returns the integers listed on the first line of
|
|
[fname]. *)
|
|
let ints_of_file fname = Aoc.string_of_file fname |> Aoc.ints_of_string
|
|
|
|
let _ =
|
|
Aoc.main ints_of_file [ (string_of_int, part 25); (string_of_int, part 75) ]
|