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)
    }