Speed2.scala

package speed2

object Speed2 {
  def main(args: Array[String]): Unit = {
    var world = World()
    world = world.newContainer()
    world = world.addWater(0, 1.0)
    world = world.newContainer()
    world = world.connectTo(0, 1)
    println(world.getAmount(0))
  }
}

case class World(containers: Vector[Container] = Vector[Container]()) {
  def newContainer(): World = {
    val newContainer = Container(containers.length, 0.0)
    World(containers :+ newContainer)
  }

  def addWater(idx: Int, amount: Double): World = {
    val container = containers(idx)
    World(
      containers.updated(idx, Container(container.nextIdx, container.amount + amount))
    )
  }

  def connectTo(idx: Int, other: Int): World = {
    val container = containers(idx)
    val otherContainer = containers(other)
    val oldNext = container.nextIdx
    World(containers.updated(idx, container.copy(nextIdx = containers(other).nextIdx)).updated(other, otherContainer.copy(nextIdx = oldNext)))
  }

  private def indicesOfGroup(head: Int): Vector[Int] = {
    @scala.annotation.tailrec
    def _indicesOfGroup(now: Int, acc: Vector[Int]): Vector[Int] = {
      val nextIdx = containers(now).nextIdx
      if (nextIdx == head) {
        acc
      } else {
        _indicesOfGroup(nextIdx, acc :+ nextIdx)
      }
    }

    _indicesOfGroup(head, Vector(head))
  }

  private def updateGroup(idx: Int): World = {
    val groups = indicesOfGroup(idx)
    val totalAmount = groups.map(idx => containers(idx).amount).sum
    val groupSize = groups.size
    val newAmount = totalAmount / groupSize

    World(
      groups.foldLeft(containers)((acc, x) => acc.updated(x, acc(x).copy(amount = newAmount)))
    )
  }

  def getAmount(idx: Int): Double = {
    val newWorld = updateGroup(idx: Int)
    newWorld.containers(idx).amount
  }
}

case class Container(nextIdx: Int, amount: Double)