Speed2.scala

package speed2

case class World(containers: Map[Symbol, Container]) {
  def getAmount(key: Symbol): (World, Double) = {
    val newWorld = updateGroup(key)
    (newWorld, newWorld.containers(key).amount)
  }

  def updateGroup(key: Symbol): World = {
    @annotation.tailrec
    def findLoop(current: Symbol, target: Symbol, acc: List[Symbol]): List[Symbol] = {
      val c = containers(current)
      if (c.next == target) current :: acc
      else findLoop(c.next, target, current :: acc)
    }

    val members = findLoop(key, key, List.empty)
    val newAmount = members.map(containers(_).amount).sum / members.size
    World(members.foldLeft(containers){ (containers, member) =>
      containers.updated(member, containers(member).copy(amount = newAmount))
    })
  }

  def addWater(key: Symbol, amount: Double): World = {
    val c = containers(key)
    World(containers.updated(key, c.copy(amount = c.amount + amount)))
  }

  def connetTo(key1: Symbol, key2: Symbol): World = {
    val next1 = containers(key1).next
    val next2 = containers(key2).next
    val newNext1 = next2
    val newNext2 = next1

    World(containers
      .updated(key1, containers(key1).copy(next = newNext1))
        .updated(key2, containers(key2).copy(next = newNext2))
    )
  }
}

case class Container(amount: Double, next: Symbol)

object Container {
  def apply(name: Symbol): Container = Container(0.0, name)
}

object Speed2 extends App {
  val world = World(Map(
    'a -> Container('a),
    'b -> Container('b),
    'c -> Container('c),
    'd -> Container('d),
  ))

  val world1 = world
    .addWater('a, 12)
    .addWater('d, 8)
    .connetTo('a, 'b)
    .connetTo('b, 'c)

  val (world2, aAmount) = world1.getAmount('a)
  val (world3, dAmount) = world2.getAmount('d)

  println((aAmount, dAmount))
}