Theme Park

ThemePark.scala

package study

object ThemePark {

  def calc(runNum: Int, maxNum: Int, group: Array[Int]):Int = {
    var itr = Stream.continually(group).flatten

    def takePeople(c: Int): Int = {
      var people = 0
      var count = 0
      for (x <- itr.take(group.size).takeWhile(x => maxNum >= people + x)) {
        people = people + x
        count = count + 1
      }
      itr = itr.drop(count)
      people
    }

    ((0 until runNum).map(_ => takePeople(maxNum)) sum)
  }
  
  
  def main(args: Array[String]) {
    val writer = new java.io.PrintWriter("a-large.out")
    try {
      process(io.Source.fromFile("C-large-practice.in").getLines)(writer.println)
    } finally {
      writer.flush()
      writer.close()
    }
  }
 
  def process(lineIn: Iterator[String])(lineOut: String => Unit) = {
    for (i <- 1 to lineIn.next().toInt) {
      val Array(runNum, maxNum, _) = lineIn.next().split(" ").map(_.toInt)
      val group = lineIn.next().split(" ").map(_.toInt)

      println(i)
 
      lineOut(s"Case #$i: ${calc(runNum, maxNum, group)}")
    }
  }
}