Scala实现一个immutable union/find set

union/find set 的接口定义通常以下 

trait UnionFindSet {
   def find(x: Int): Int //return the root of x
   def union(x: Int, y: Int): Union   //put x and y into the same set
}

在算法课程里，（使用C/Java)有一个经典的实现，就是利用一个足够大的数组，其中每一个位置对应一个元素，值对应该元素父节点的位置。具体的细节通常都比较清楚，就不在这里赘述了。

固然，scala里面也是有数组的，也能够用一样的方式实现，但就少了一点乐趣。我想实现的union/find set具备如下两个特色：

1. immutable，这个符合scala的函数式精神；

2. 能够在非Int元素的状况下使用；

代码以下：

/**
 * Created by senyuanwang on 14-10-6.
 */
class UnionFindSet[A: ID](val map: Map[Int, Int]) {
  private val ev = implicitly[ID[A]]
  var internalMap = map.withDefault(x => x)

  def this() = this(Map.empty)

  def find(a: A): Int = find(ev.id(a))

  def ?(a: A) = find(a)

  def union(x: A, y: A): UnionFindSet[A] = union(ev.id(x), ev.id(y))

  def union(x: Int, y: Int): UnionFindSet[A] = {
    val px = find(x)
    val py = find(y)
    new UnionFindSet[A](internalMap + (px -> py))
  }

  def ?(p: (A, A)) = {
    val px = find(p._1)
    val py = find(p._2)
    px == py
  }

  def +(p: (A, A)) = union(p._1, p._2)

  private def find(x: Int): Int = {
    val p = internalMap(x)
    if (p == x) {
      p
    } else {
      internalMap += (x -> find(p))
      internalMap(x)
    }
  }
}


trait ID[A] {
  def id(a: A): Int
}

object UnionFindSet {
  implicit def intID = new ID[Int] {
    def id(a: Int) = a
  }

  def apply[A: ID]() = new UnionFindSet[A]()
}

UnionFindSet内部借组了一个var internalMap: Map[Int, Int]。由于在调用find的时候，（为了性能考虑）使用了路径压缩，因此内部存储的结构必然会改变。但对外应该是透明的。每次union都返回一个新的UnionFindSet，固然新的要基于老的进行建立，须要将保存内部状态的map进行传递。由于Map也是immutable的，因此不用担忧新的set的更新会影响到老的map。

如下是一段使用UnionFindSet实现Krushkal算法的例子：

def kruskal(es: List[(Int, Int, Int)], uf: UnionFindSet[Int], sum: Int): Int =
    es match {
      case Nil => sum
      case (x, y, d) :: tail if (uf ? (x -> y)) => kruskal(tail, uf, sum)
      case (x, y, d) :: tail => kruskal(tail, uf + (x -> y), sum + d)
    }