Rotate.scala

object Rotate extends App {

  def whoWin(k: Int, board: IndexedSeq[String]): String = {

    val transformed = {
      def gravity(line: String): String = {
        val (space, block) = line partition (_ == '.')
        space ++ block
      }
      board map gravity
    }

    def win(color: Char): Boolean = {
      val directions = Seq((1, 1), (1, 0), (1, -1), (0, 1))

      def hasSeq(i: Int, j: Int, di: Int, dj: Int): Boolean = {
        val (i1, j1) = (i + (k - 1) * di, j + (k - 1) * dj)
        0 <= i && i1 < transformed.size && 0 <= j1 && j1 < transformed.size
      }

      (for {
        (line, i) <- transformed.zipWithIndex.toStream
        (c, j) <- line.zipWithIndex if c == color
        (di, dj) <- directions if hasSeq(i, j, di, dj)
        lines = for {
          l <- 1 until k
        } yield transformed(i + di * l)(j + dj * l)
        if lines.forall { _ == color }
      } yield lines).nonEmpty
    }

    (win('R'), win('B')) match {
      case (false, false) => "Neither"
      case (false, true)  => "Blue"
      case (true, false)  => "Red"
      case (true, true)   => "Both"
    }

  }

  def process(lineIn: Iterator[String])(lineOut: String => Unit) =
    for (i <- 1 to lineIn.next().toInt) {
      val Array(n, k) = lineIn.next() split ' ' map (_.toInt)
      val board = Vector.fill(n)(lineIn.next())
      lineOut(s"Case #$i: ${whoWin(k, board)}")
    }

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

}

RotateTest.scala

import org.scalatest._

import Rotate._

class RotateTest extends FunSuite {

  test("sample #1") {
    val board = """.......
.......
.......
...R...
...BB..
..BRB..
.RRBR..""".lines.toVector

    assert(whoWin(3, board) === "Neither")
  }
  
  
  test("sample #2") {
    val board = """......
......
.R...R
.R..BB
.R.RBR
RB.BBB""".lines.toVector


    assert(whoWin(4, board) === "Both")
  }

  test("sample #3") {
    val board = """R...
BR..
BR..
BR..""".lines.toVector


    assert(whoWin(4, board) === "Red")
  }

  test("sample #4") {
    val board = """B..
RB.
RB.""".lines.toVector


    assert(whoWin(3, board) === "Blue")
  }

  test("sample case") {
    val input = """4
7 3
.......
.......
.......
...R...
...BB..
..BRB..
.RRBR..
6 4
......
......
.R...R
.R..BB
.R.RBR
RB.BBB
4 4
R...
BR..
BR..
BR..
3 3
B..
RB.
RB.""".lines

    val expected = """Case #1: Neither
Case #2: Both
Case #3: Red
Case #4: Blue""".lines

    lineComparison(input, expected)
  }

  test("full small case") {
    val input = io.Source.fromFile("A-small-practice.in").getLines()
    val expected = io.Source.fromFile("a.small.out.ref").getLines()

    lineComparison(input, expected)
  }

  test("full large case") {
    val input = io.Source.fromFile("A-large-practice.in").getLines()
    val expected = io.Source.fromFile("a.large.out.ref").getLines()

    lineComparison(input, expected)
  }

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

}