Matthew Gretton-Dann
2159a5fc5e
We notice that we're repeating calculations at each step, so use a map to ensure we do each stone ID once per step.
54 lines
2.2 KiB
OCaml
54 lines
2.2 KiB
OCaml
module IntMap = Map.Make (Int)
|
|
|
|
let load_file fname =
|
|
match In_channel.with_open_text fname In_channel.input_line with
|
|
| Some x -> x
|
|
| None -> failwith "load_file"
|
|
|
|
let rec apply_n n fn arg = if n <= 0 then arg else apply_n (n - 1) fn (fn arg)
|
|
let update_count n = function None -> Some n | Some x -> Some (x + n)
|
|
|
|
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
|
|
|
|
(** [calc_blink_rec acc lst] returns an updated map based off [acc] with the
|
|
result of apply a blink step to the stones in [lst]. Entries in [lst] are
|
|
pairs of [(stone id, count)]. [acc] and the resulting map have keys which
|
|
are stone id, and values which are count. *)
|
|
let rec calc_blink_rec acc = function
|
|
| [] -> acc
|
|
| (0, n) :: t -> calc_blink_rec (IntMap.update 1 (update_count n) acc) t
|
|
| (x, n) :: t 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
|
|
calc_blink_rec acc t
|
|
| (x, n) :: t ->
|
|
calc_blink_rec (IntMap.update (x * 2024) (update_count n) acc) t
|
|
|
|
(** [calc_blink map] calculates how a collection of stones changes in a blink.
|
|
[map] is a map with key of stone ID, and value number of times that stone
|
|
appears. The result is a map with similar key, value pairs.
|
|
|
|
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.to_list map |> calc_blink_rec IntMap.empty
|
|
|
|
(** [part n str] returns the number of stones after [n] blinks, given an initial
|
|
string [str] of space seperated stone IDs. *)
|
|
let part n str =
|
|
let map = Aoc.ints_of_string str |> map_of_ints |> apply_n n calc_blink in
|
|
IntMap.fold (fun _ v acc -> v + acc) map 0
|
|
|
|
let _ =
|
|
Aoc.main load_file [ (string_of_int, part 25); (string_of_int, part 75) ]
|