advent-of-code-2020/day11.ml

(* Common util. Should probably put this in a module or something.. *)
let get_one_line file =
  try Some (input_line file) with
    End_of_file -> None

let get_lines file =
  let rec input lines =
    match get_one_line file with
      Some line -> input (line :: lines)
    | None -> List.rev lines
  in input []

let read_file file =
  let channel = open_in(file) in
  get_lines channel

let explode s = (* convert a string to a list of chars *)
  let rec exp i l =
   if i < 0 then l else exp (i - 1) (s.[i] :: l) in
  exp (String.length s - 1) [];;

let print_grid l =
  print_newline ();
  List.iter (fun line ->
    List.iter (Printf.printf "%c") line;
    Printf.printf "\n") l

let (<<) f g x = f(g(x))

(* --- *)

let get_value l (x, y) = (* if in bounds, Some of the value, else None *)
  try let row = List.nth l x in
    Some (List.nth row y)
  with
  | Failure _ -> None
  | Invalid_argument _ -> None

let dirs = [(-1, -1); (-1, 0); (-1, 1); (0, -1); (0, 1); (1, -1); (1, 0); (1, 1)]

let get_neighbors (x, y) l = (* returns list of Options. None means off the grid *)
  let coords = List.map (fun (dx, dy) -> (x + dx, y + dy)) dirs in
  List.map (get_value l) coords

let rec walk (x, y) l (dx, dy) = (* walk from (x, y) in direction (dx, dy) until you hit something or fall off *)
  let coord = (x + dx, y + dy) in
  match get_value l coord with
  | Some '.' -> walk coord l (dx, dy)
  | a -> a

let get_far_neighbors (x, y) l = (* returns list of Options. None means off the grid *)
  List.map (walk (x, y) l) dirs

let num_occupied = (* count the number of # in a List of char options *)
  List.length << List.filter (function
    | None -> false
    | Some v -> Char.equal v '#')

let change_seat is_far (x, y) l = (* is_far is for the 2nd problem *)
  let neighbors_fun = if is_far then get_far_neighbors else get_neighbors in
  let max_neighbors = if is_far then 5 else 4 in

  let neighbors = neighbors_fun (x, y) l in (* make a list of all neighbors, either near or far *)
  match get_value l (x, y) with
    | None -> '?' (* eh? *)
    | Some 'L' -> if num_occupied neighbors = 0 then '#' else 'L'
    | Some '#' -> if num_occupied neighbors >= max_neighbors then 'L' else '#'
    | Some s -> s

let change_seats is_far l =
  List.mapi (fun x ->
    List.mapi (fun y _ ->
      change_seat is_far (x, y) l
    )
  ) l

let rec start is_far grid = (* keep changing seats until the grid stops changing *)
  (* print_grid grid; *)
  let grid' = change_seats is_far grid in
  if grid = grid'
  then grid'
  else start is_far grid'

let () =
  let grid = read_file "day11.txt" |> List.map explode in

  start false grid (* false means near-neighbors *)
    |> List.flatten
    |> List.filter (Char.equal '#') |> List.length (* count occupied seats *)
    |> Printf.printf "near-neighbor count %i\n";

  start true grid
    |> List.flatten
    |> List.filter (Char.equal '#') |> List.length
    |> Printf.printf "far-neighbor count %i\n";