Connected Component Problem

cc.scala

package me.invkrh.train

import org.apache.spark.{SparkContext, SparkConf}
import org.apache.spark.rdd.RDD


/**
 * Created with IntelliJ IDEA.
 * User: invkrh
 * Date: 22/07/15
 * Time: 23:14
 */

object ConnectedComponent extends App {

  val conf = new SparkConf()
    .setAppName(this.getClass.getName)
    .setMaster("local[*]")

  val sc = new SparkContext(conf)

  val arcs = sc.makeRDD(Seq((1, 2), (2, 3), (2, 4), (5, 6)))

  val input = arcs.flatMap {
    case pair => Array(pair, pair.swap)
  }.map {
    case (left, right) => (Set(left), Set(right))
  }

  val output = connectedComponent(input.persist()).keys.collect foreach println

  def connectedComponent(dataSet: RDD[(Set[Int], Set[Int])]): RDD[(Set[Int], Set[Int])] = {
    val keys = dataSet.keys
    val cnt = keys.count()
    val distCnt = keys.distinct().count()

    if (cnt != distCnt) {
      val next = dataSet.reduceByKey(_ union _)
        .flatMap { case (cc, neighbours) =>
        if (neighbours.nonEmpty) {
          neighbours.map(n => (cc + n, neighbours - n))
        } else {
          Array((cc, Set[Int]()))
        }
      }.persist()
      next.count()
      dataSet.unpersist(false)
      connectedComponent(next)
    } else {
      dataSet
    }
  }

}