Speed3.scala

package speed3

object Speed3 {
  def main(args: Array[String]): Unit = {
    var world = World(Map(
      "a" -> Container("a"),
      "b" -> Container("b"),
      "c" -> Container("c"),
      "d" -> Container("d")
    ))

    val w1 = world
      .addWater("a", 12)
      .addWater("d", 8)
      .connectTo("a", "b")
      .connectTo("b", "c")

    val (w2, a) = w1.getAmount("a")
    val (w3, d) = w2.getAmount("d")

    println((a, d))
    println(w3)
  }
}

case class World(containers: Map[String, Container]) {
  def findRootAndCompress(key: String): (World, String) = {
    val c = containers(key)
    val parentKey = c.parentKey
    if (parentKey != key) {
      val (newWorld, newParentKey) = findRootAndCompress(parentKey)
      (World(newWorld.containers.updated(key, c.copy(parentKey = newParentKey))), newParentKey)
    } else (this, parentKey)
  }

  def getAmount(key: String): (World, Double) = {
    val (w, parentKey) = findRootAndCompress(key)
    (w, w.containers(parentKey).amount)
  }

  def addWater(key: String, amount: Double): World = {
    val (w, rootKey) = findRootAndCompress(key)
    val root = w.containers(rootKey)
    val newAmount = root.amount + amount / root.size
    World(w.containers.updated(rootKey, root.copy(amount = newAmount)))
  }

  def connectTo(thisKey: String, otherKey: String): World = {
    val (world1, rootKey1) = findRootAndCompress(thisKey)
    val (world2, rootKey2) = world1.findRootAndCompress(otherKey)
    if (rootKey1 == rootKey2) world2
    else {
      val root1 = world2.containers(rootKey1)
      val root2 = world2.containers(rootKey2)
      val size1 = root1.size
      val size2 = root2.size
      val newSize = size1 + size2
      val newAmount = (root1.amount * size1 + root2.amount * size2) / newSize
      if (size1 <= size2) {
        World(world2.containers
          .updated(rootKey1, root1.copy(parentKey = rootKey2))
          .updated(rootKey2, root2.copy(amount = newAmount, size = newSize))
        )
      } else {
        World(world2.containers
          .updated(rootKey2, root2.copy(parentKey = rootKey1))
          .updated(rootKey1, root1.copy(amount = newAmount, size = newSize))
        )
      }
    }
  }
}

case class Container(parentKey: String, amount: Double, size: Int)

object Container {
  def apply(key: String): Container = new Container(key, 0.0, 1)
}