ManageYourEnergy.scala

object ManageYourEnergy extends App {

  def maximumGain(e: Int, r: Int, vs: Seq[Int]): Int = {

    @annotation.tailrec
    def maximumGain0(current: Int, vs: List[Int], acc: Int = 0): Int = vs match {
      case Nil => acc
      case x :: xs =>
        val consumable = math.max(xs.take(e / r).zipWithIndex.find(_._1 > x) match {
          case None => current
          case Some((y, i)) => current - (e - r * (i + 1))
        }, 0)
        maximumGain0(current - consumable + r, xs, acc + x * consumable)
    }

    maximumGain0(e, vs.toList)

  }

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

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

}

ManageYourEnergyTest.scala

import org.scalatest._

import ManageYourEnergy._

class ManageYourEnergyTest extends FunSuite {

  test("sample #1") {
    assert(maximumGain(5, 2, List(2,1)) === 12)
  }
  

  test("sample #2") {
    assert(maximumGain(5, 2, List(1,2)) === 12)
  }
  
  test("sample #3") {
    assert(maximumGain(3, 3, List(4,1,3,5)) === 39)
  }
  
  test("small #5") {
    assert(maximumGain(5, 1, List(3, 1, 2)) === 19)
  }
  
  
  test("sample case") {
    val input = """3
5 2 2
2 1
5 2 2
1 2
3 3 4
4 1 3 5""".lines

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

    process(input) { s =>
      for (line <- s.lines)
        assert(line === expected.next())
    }
  }

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

    process(input) { s =>
      for (line <- s.lines)
        assert(line.trim === expected.next().trim)
    }
  }

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

    process(input) { s =>
      for (line <- s.lines)
        assert(line.trim === expected.next().trim)
    }
  }
}