Add union by rank to UnionFind

Author: sim642

src/main/scala/eu/sim642/adventofcodelib/UnionFind.scala

@@ -1,34 +1,49 @@
 package eu.sim642.adventofcodelib
 
+import eu.sim642.adventofcodelib.UnionFind.Node
+
 import scala.annotation.tailrec
 
-class UnionFind[A](val reprs: Map[A, A]) {
-  // TODO: optimize
+class UnionFind[A](val nodes: Map[A, Node[A]]) {
 
-  def this(items: Seq[A]) = {
-    this(items.map(x => x -> x).toMap)
+  def this(xs: Seq[A]) = {
+    this(xs.map(x => x -> Node(x, 1)).toMap)
   }
 
   @tailrec
-  final def findRepr(x: A): A = {
-    val repr = reprs(x)
-    if (x == repr)
-      repr
+  final def findReprNode(x: A): Node[A] = {
+    // TODO: path compression
+    val node = nodes(x)
+    if (x == node.parent)
+      node
     else
-      findRepr(repr)
+      findReprNode(node.parent)
   }
 
+  def findRepr(x: A): A = findReprNode(x).parent
+
   def sameRepr(x: A, y: A): Boolean =
     findRepr(x) == findRepr(y)
 
   def unioned(x: A, y: A): UnionFind[A] = {
-    val xRepr = findRepr(x)
-    val yRepr = findRepr(y)
-    new UnionFind(reprs + (yRepr -> xRepr))
+    val xNode = findReprNode(x)
+    val yNode = findReprNode(y)
+    val xRepr = xNode.parent
+    val yRepr = yNode.parent
+    if (xRepr == yRepr) // sameRepr inlined
+      this
+    else if (xNode.size >= yNode.size)
+      new UnionFind(nodes + (xRepr -> xNode.copy(size = xNode.size + yNode.size)) + (yRepr -> yNode.copy(parent = xRepr)))
+    else
+      new UnionFind(nodes + (yRepr -> yNode.copy(size = xNode.size + yNode.size)) + (xRepr -> xNode.copy(parent = yRepr)))
   }
 
   def groups(): Seq[Seq[A]] =
-    reprs.keys.groupBy(findRepr).values.map(_.toSeq).toSeq
+    nodes.keys.groupBy(findRepr).values.map(_.toSeq).toSeq
+
+  override def toString: String = nodes.toString()
+}
 
-  override def toString: String = reprs.toString()
+object UnionFind {
+  case class Node[A](parent: A, size: Int)
 }