BFFs.scala

object BFFs extends App {

  def maxCircleSize(bffs: IndexedSeq[Int]): Int = {

    val (clusters, membership) = clusterize(bffs)

    println((clusters, membership))

    val maxLoopSize = clusters.map(_.size).max 

    val coupleCircleSize = clusters.zipWithIndex.filter(_._1.size == 2).map{
      case (cluster, index) => coupleClusterSize(bffs, cluster, membership.zipWithIndex.filter(_ == index).map(_._2).toSet)
    }.sorted.take(2).sum

    maxLoopSize max coupleCircleSize
  }

  def clusterize(bffs0: IndexedSeq[Int]): (IndexedSeq[Set[Int]], IndexedSeq[Int]) = {
    // bffs0 = IndexedSeq(2,3,4,1)
    val bffs = bffs0.map(_ -1)

    println(s"bffs ==> $bffs")
    // bffs = IndexedSeq(1,2,3,0)

    def findCluster(notClustered: Set[Int], clustered: Map[Int, Int], clusters: IndexedSeq[Set[Int]]): (IndexedSeq[Set[Int]], IndexedSeq[Int]) = {
      //Thread.sleep(300)
      //println(s"===> ${(notClustered, clustered, clusters)}")

      if (notClustered.isEmpty) (clusters, clustered.toIndexedSeq.sorted.map(_._2))
      else {
        def search(current: Int, visited: Map[Int, Int], count: Int): (Seq[Int], Either[Int, Int]) = {
          //println((current, visited, count))
          val next = bffs(current)
          if (clustered contains next) (visited.toSeq.sortBy(_._2).map(_._1) :+ current, Left(clustered(next)))
          else if (visited contains next) (visited.toSeq.sortBy(_._2).map(_._1) :+ current, Right(count - visited(next) + 1))
          else search(next, visited + ((current -> count)), count + 1)
        }
        search(notClustered.head, Map.empty[Int, Int], 0) match {
          case (members, Left(clusterIndex)) =>
            findCluster(notClustered -- members, clustered ++ (members.map(_ -> clusterIndex)), clusters)
          case (members, Right(loopSize)) =>
            findCluster(notClustered -- members, clustered ++ (members.map(_ -> clusters.size)), clusters :+ members.takeRight(loopSize).toSet)
        }
      }
    }

    findCluster(bffs.toSet, Map.empty[Int, Int], IndexedSeq.empty[Set[Int]])
  }

  def coupleClusterSize(bffs0: IndexedSeq[Int], couple: Set[Int], members: Set[Int]): Int = {
    val bffs = bffs0.map(_ -1)

    val g = (Map.empty[Int, Set[Int]] /: bffs.zipWithIndex) { case (graph, (bff, index)) =>
      graph + ((bff -> (graph.getOrElse(bff, Set.empty[Int]) + index)))
    }

    //println(g)

    def height(index: Int): Int = {
      val childHeights = {
        for {
          children <- g.getOrElse(index, Set.empty[Int]) if !(couple contains children)
        } yield height(children)
      }
      if (childHeights.isEmpty) 0 else childHeights.max + 1
    }
    2 + (couple map height).sum
  }

  def process(lineIn: Iterator[String])(lineOut: String => Unit) =
    for (i <- 1 to lineIn.next().toInt) {
      lineIn.next()
      lineOut(s"Case #$i: ${maxCircleSize(lineIn.next().split(" ").map(_.toInt))}")
    }

  val filename = "C-small-practice"
  val writer = new java.io.PrintWriter(filename + ".out")
  try {
    process(io.Source.fromFile(filename + ".in").getLines) { s =>
      writer.println(s); writer.flush()
    }
  } finally {
    writer.flush(); writer.close()
  }
}

BFFsTest.scala

import org.scalatest._

import BFFs._

class BFFsTest extends FunSuite with Matchers {

  ignore("clusterize") {
    assert(clusterize(IndexedSeq(2,3,4,1)) === (IndexedSeq(Set(0, 1, 2, 3)), IndexedSeq(0, 0, 0, 0)))
  }

  ignore("coupleClusterSize") {
    assert(coupleClusterSize(IndexedSeq(3,3,4,3), Set(2,3), Set(0, 1, 2, 3)) === 3)
  }

  ignore("sample #1") {
    assert(maxCircleSize(IndexedSeq(2,3,4,1)) === 4)
  }

  ignore("sample #2") {
    assert(maxCircleSize(IndexedSeq(3,3,4,1)) === 3)
  }

  ignore("sample #3") {
    assert(maxCircleSize(IndexedSeq(3,3,4,3)) === 3)
  }

  ignore("sample #4") {
    assert(maxCircleSize(IndexedSeq(7,8,10,10,9,2,9,6,3,3)) === 6)
  }

  ignore("case #5") {
    assert(maxCircleSize(IndexedSeq(2, 3, 1, 10, 3, 5, 6, 7, 8, 9)) === 3)
  }

  test("case #18") {
    assert(maxCircleSize(IndexedSeq(3, 4, 1, 3, 6, 1)) === 6)
  }

  ignore("sample case") {
    val input = """4
4
2 3 4 1
4
3 3 4 1
4
3 3 4 3
10
7 8 10 10 9 2 9 6 3 3""".lines

    val expected = """Case #1: 4
Case #2: 3
Case #3: 3
Case #4: 6""".lines

    lineComparison(input, expected)
  }

  ignore("full small case") {
    val input = io.Source.fromFile("B-small-practice.in").getLines()
    val expected = io.Source.fromFile("B-small-practice.out").getLines()

    lineComparison(input, expected)
  }

  ignore("full large case") {
    val input = io.Source.fromFile("B-large-practice.in").getLines()
    val expected = io.Source.fromFile("B-large-practice.out").getLines()

    lineComparison(input, expected)
  }

  def lineComparison(input: Iterator[String], expected: Iterator[String]): Unit = {
    process(input) { s =>
      for (line <- s.lines) assert(line.trim === expected.next().trim)
    }
    assert(expected.hasNext === false, "Finished too fast.")
  }

}