Rotate.scala

object Rotate extends App {

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

    val transformed =
      for (line <- board; (space, block) = line partition (_ == '.'))
        yield space ++ block

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

      def ?(x: Int, d: Int): Boolean = {
        val y = x + (k - 1) * d
        0 <= y && y < transformed.size
      }

      if ((for {
        (line, i) <- transformed.zipWithIndex.toStream
        (c, j) <- line.zipWithIndex if c == color
        (di, dj) <- directions if ?(i, di) && ?(j, dj)
        lines = 1 until k map { l => transformed(i + di * l)(j + dj * l) }
        if lines forall { _ == color }
      } yield lines).isEmpty) 0 else 1
    }

    Vector("Neither", "Blue", "Red", "Both")(2 * win('R') + win('B'))

  }

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

}