Advent of Code 14.scala

import java.io.FileInputStream
import scala.io.StdIn

@main def solve14() =

  sealed trait Command
  case class ValueCommand(address: Int, value: Int) extends Command
  case class MaskCommand(mask: String) extends Command
  
  def solve1(commands: Vector[Command], mask: String = "X" * 36, acc: Map[Int, String] = Map()): BigInt = {
    commands.headOption match {
      case Some(ValueCommand(address, value)) =>
        val binaryString = value.toBinaryString
        val filledBinaryString = ("0" * (36 - binaryString.length)) + binaryString
        val maskedString = (filledBinaryString zip mask).map { (c0, c1) =>
          (c0, c1) match {
            case (_, '0') =>
              '0'
            case (_, '1') =>
              '1'
            case (v, 'X') =>
              v
          }
        }.mkString
        val nextAcc = acc.updated(address, maskedString)
        solve1(commands.tail, mask, nextAcc)
      case Some(MaskCommand(mask)) =>
        solve1(commands.tail, mask, acc)
      case _ =>
        acc.values.map(BigInt(_, 2)).sum
    }
  }
  
  def floatString(input: Vector[Vector[Char]], acc: Vector[Vector[Char]] = Vector(Vector())): Vector[Vector[Char]] = {
    input.headOption match {
      case Some(chars) =>
        floatString(input.tail, acc.flatMap(c => chars.map(c :+ _)))
      case _ =>
        acc
    }
  }
  
  def solve2(commands: Vector[Command], mask: String = "X" * 36, acc: Map[BigInt, BigInt] = Map()): BigInt = {
    commands.headOption match {
      case Some(ValueCommand(address, value)) =>
        val binaryString = address.toBinaryString
        val filledBinaryString = ("0" * (36 - binaryString.length)) + binaryString
        val possibleValues = (filledBinaryString zip mask).map { (c0, c1) =>
          (c0, c1) match {
            case (v, '0') =>
              Vector(v)
            case (_, '1') =>
              Vector('1')
            case (v, 'X') =>
              Vector('0', '1')
          }
        }.toVector
        val nextAcc = floatString(possibleValues).foldLeft(acc)((acc, address) => _.updated(BigInt(address.mkString, 2), value))
        solve2(commands.tail, mask, nextAcc)
      case Some(MaskCommand(mask)) =>
        solve2(commands.tail, mask, acc)
      case _ =>
        acc.values.sum
    }
  }

  val in = new FileInputStream("example14-1.in")
  System.setIn(in)

  val inputs = Iterator.continually(StdIn.readLine())
  val commands = inputs.takeWhile(_ != null).toVector.map { line =>
    if line.startsWith("mask") then
      MaskCommand(line.filter(c => c.isDigit || c == 'X'))
    else
      val Array(address, value) = line.filter(c => c.isDigit || c == ' ').split(" ").filter(_.nonEmpty)
      ValueCommand(address.toInt, value.toInt)
  }
  val answer1 = solve1(commands)
  println(answer1)
  val answer2 = solve2(commands)
  println(answer2)