10.koto

# https://adventofcode.com/2020/day/10

from io import extend_path, read_to_string
from koto import script_dir

parse_input = |input|
  input
    .lines()
    .each |line| line.to_number()
    .to_list()
    .sort()

joltage_differences = |adapters|
  differences = [0, 0, 0, 0]
  previous = 0
  for adapter in adapters
    difference = adapter - previous
    differences[difference] += 1
    previous = adapter
  differences[3] += 1
  differences

count_arrangements = |adapters|
  # The first adapter has a single arrangement
  arrangements = [1]
  # Iterate over the remaining adapters
  for i in 1..adapters.size()
    adapter = adapters[i]
    # Adapters with ratings <= 3 are also valid start points
    arrangements_for_adapter = if adapter <= 3 then 1 else 0
    # Add the number of arrangements from previous adapters within range
    for j in i - 1..=0
      if adapter - adapters[j] <= 3
        arrangements_for_adapter += arrangements[j]
      else
        break
    arrangements.push arrangements_for_adapter
  # The result is the arrangement count for the last adapter
  arrangements.last()

@main = ||
  adapters = io.extend_path script_dir, "input", "10"
    >> read_to_string
    >> parse_input
  differences = joltage_differences adapters
  print "Part one: ${differences[1] * differences[3]}" # 2040
  print "Part two: ${count_arrangements adapters}" # 28346956187648

@tests =
  @pre_test: ||
    self.adapters_1 =
      parse_input read_to_string extend_path script_dir, "input", "10-example-1"
    self.adapters_2 =
      parse_input read_to_string extend_path script_dir, "input", "10-example-2"

  @test part_one: ||
    differences = joltage_differences self.adapters_1
    assert_eq differences[1], 7
    assert_eq differences[3], 5

    differences = joltage_differences self.adapters_2
    assert_eq differences[1], 22
    assert_eq differences[3], 10

  @test part_two: ||
    assert_eq (count_arrangements self.adapters_1), 8
    assert_eq (count_arrangements self.adapters_2), 19208