Day10.scala

package sungkmi.aoc2020.day10

def parse(inputString: String): List[Int] =
  inputString.split("\n").map(BigInt(_).toInt).toList

def countJolts(joltageRatings: List[Int]): (Int, Int, Int) =
  val sorted = 0 :: joltageRatings.sorted
  sorted.zip(sorted.tail).foldLeft((0, 0, 1)):
    case ((j1, j2, j3), (i1, i2)) =>
      (i2 - i1) match
        case 1 => (j1 + 1, j2, j3)
        case 2 => (j1, j2 + 1, j3)
        case 3 => (j1, j2, j3 + 1)
        case _ => throw new Exception("Wrong joltage diff")

def countWays(joltageRatings: List[Int]): BigInt =
  (0 :: joltageRatings.sorted).foldRight(List.empty[(Int, BigInt)]):
    case (n, Nil) => (n, BigInt(1)) :: Nil
    case (n, xs) => (n, xs.takeWhile(_._1 <= (n + 3)).map(_._2).sum) :: xs
  .head._2

@main def part1: Unit =
  val (j1, j2, j3) = countJolts(joltageRatings)
  val ans = j1 * j3
  println(ans)

@main def part2: Unit =
  val ans = countWays(joltageRatings)
  println(ans)

lazy val joltageRatings: List[Int] = parse(input)

//lazy val input: String = """160

Day10Test.scala

package sungkmi.aoc2020.day10

class Day10Test extends munit.FunSuite {

  val smallInput = """16
10
15
5
1
11
7
19
6
12
4"""

  val largeInput = """28
33
18
42
31
14
46
20
48
47
24
23
49
45
19
38
39
11
1
32
25
35
8
17
7
9
4
2
34
10
3"""

  test("countJolts small input") {
    assertEquals(countJolts(parse(smallInput)), (7, 0, 5))
  }

  test("countJolts large input") {
    assertEquals(countJolts(parse(largeInput)), (22, 0, 10))
  }

  test("countWays small input") {
    assertEquals(countWays(parse(smallInput)), BigInt(8))
  }
  
  test("countWays large input") {
    assertEquals(countWays(parse(largeInput)), BigInt(19208))
  }
}