Haircut.scala

object Haircut extends App {

  def findBarber(n: Long, ms: Seq[Long]): Long = {

    def numberOfServed(time: Long): Long = ms.map(time.toLong / _).sum

    def findWatingTime(l: Long, u: Long): Long = {

      if (l == u) l
      else {
        val half = (l + u) / 2
        val served = numberOfServed(half)
        val servedL = numberOfServed(l)
        val servedU = numberOfServed(u)

        if (servedL == served) l
        else if (servedU == served) half
        else if (n <= served) findWatingTime(l, half)
        else findWatingTime(half, u)
      }
    }
    val t = findWatingTime(0, n * ms.max).toInt
    val waitingTime = Stream.from(t, -1).takeWhile(n => numberOfServed(n) == numberOfServed(t)).last

    val candidate = for {
      (m, i) <- ms.zipWithIndex.toVector if waitingTime % m == 0
    } yield i + 1
    val alreadyServed = numberOfServed(waitingTime)
    candidate((n - alreadyServed - 1).toInt)
  }

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

  val filename = "B-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()
  }

}

HaircutTest.scala

import org.scalatest._

import Haircut._

class HaircutTest extends FunSuite {

  test("sample #1") {
    assert(findBarber(4, Seq(10, 5)) === 1)
  }

  test("sample #2") {
    assert(findBarber(12, Seq(7, 7, 7)) === 3)
  }

  test("sample #3") {
    assert(findBarber(8, Seq(4, 2, 1)) === 1)
  }
  
  test("sample case") {
    val input = """3
2 4
10 5
3 12
7 7 7
3 8
4 2 1""".lines

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

    lineComparison(input, expected)
  }

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

    lineComparison(input, expected)
  }

  ignore("full large case") {
    val input = io.Source.fromFile("D-large-practice.in").getLines()
    val expected = io.Source.fromFile("D-large-practice.out").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.")
  }

}