ex3-1.scala

object Speed1 {
  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.containers)
  }
}

case class World(containers: Vector[Container] = Vector[Container]()) {

  def newContainer(): World = {
    val newContainer = Container(containers.length, Group(Set(containers.length), 0.0))
    World(containers :+ newContainer)
  }

  def addWater(containerIdx: Int, amount: Double) = {
    val group = containers(containerIdx).group
    val newGroup = group.copy(amountPerContainer = group.amountPerContainer + amount / group.members.size)
    World(
      containers.zipWithIndex.map { case (container, idx) =>
        if (group.members.contains(idx)) {
          container.copy(group = newGroup)
        } else {
          container
        }
      }
    )
  }

  def connectTo(idx1: Int, idx2: Int): World = {
    val group1 = containers(idx1).group
    val group2 = containers(idx2).group
    if (group1 == group2) return this
    val size1 = group1.members.size
    val size2 = group2.members.size
    val tot1 = group1.amountPerContainer * size1
    val tot2 = group2.amountPerContainer * size2
    val newAmount = (tot1 + tot2) / (size1 + size2)

    val newGroup: Group = Group(group1.members ++ group2.members, newAmount)
    World(
      containers.zipWithIndex.map { case (container, idx) =>
        if (newGroup.members.contains(idx)) {
          container.copy(group = newGroup)
        } else {
          container
        }
      }
    )
  }
}

case class Group(members: Set[Int], amountPerContainer: Double) {

}

case class Container(idx: Int, group: Group) {
  def getAmount: Double = group.amountPerContainer
}