pouring water scala

pouring.scala

import collection.mutable.HashSet
import collection.mutable.Queue
import collection.mutable.ArraySeq
import scala.io.Source
import java.io.File

case class Container(capacity: Int, filled: Int) {

    // overload + and - methods
    def +(toAdd: Int) = {
        assert (filled + toAdd <= capacity, { println("Too much water in container")})
        new Container(capacity, filled + toAdd)
    }

    def -(toSub: Int) = { 
        assert (filled - toSub >= 0, { println("Not enough water in container")})
        new Container(capacity, filled - toSub)
    }

    // how much can we pour into an other container ?
    def howMuchToPour(other: Container) = 
        math.min(filled, other.capacity - other.filled)

    // just print the filling level when printing a Container
    override def toString(): String = filled.toString;
}

case class State(containers : ArraySeq[Container]) {
    var parent:State = _
    
    // returns all the ancesters of a State
    def ancesters() : List[State] = 
        if (Option(parent) != None) parent :: parent.ancesters
        else                        List[State]()
    
    // returns a new State with i-th container poured into j-th container
    // no check if it's possible or not
    def pour(i: Int, j: Int, q: Int) = {
        val new_containers = containers.slice(0, containers.length) // make a copy
        new_containers(i) = containers(i) - q
        new_containers(j) = containers(j) + q
        val ns = State(new_containers)
        ns.parent = this
        ns
    }
    
    override def toString(): String = containers map{_ toString} mkString "\t"

    def next_states() = 
        for { i <- (0 until containers.length);
              j <- (0 until containers.length)
              q = containers(i) howMuchToPour containers(j)
              if (i !=j) && q > 0 }
            yield pour(i, j, q)
}

object Pouring {
    def _bfs(start: State, condition: (State) => Boolean) : List[State]  = {
        val q = Queue[State]()
        val seen =  HashSet[State]()
        q    += start    
        seen += start

        while (q.length > 0) {
            val s = q.dequeue()
            
            if (condition(s)) 
                return s :: s.ancesters
            
            s.next_states().foreach(ns => {
                if (!seen.contains(ns)) {
                    q    += ns
                    seen += ns
                }
            })
        }
        List[State]()
    }
    
    def solve(problem:Seq[Array[Int]]) {        
        val state = State(problem(1).zip(problem(2)).map(x => Container(x._1, x._2)))

        val cond_funcs = for { (filled, index) <- problem(3).zipWithIndex
                                                            .filter(x => x._1 != 0) } 
                            yield (x:State) => x.containers(index).filled == filled

        _bfs(state, (x) => cond_funcs.forall(_(x)))
            .reverse
            .foreach(s => println(s))
    }
    
    // convenient way to ask user for input
    def askForInput(question: String) = {
        println(question)
        Console.readLine
    }
    
    // parse a line into an array of int
    def line_to_int(line: String) = line.split(" ").map(_.toInt)
    
    // is a valid problem ? nb_container, capacities, start, goal
    def isProblem(problem: Seq[Array[Int]]) = 
        problem.length == 4 &&
        problem(0)(0) > 0 && 
        problem(1).length == problem(0)(0) && 
        problem(2).length == problem(0)(0) && 
        problem(3).length == problem(0)(0)

    def usage() = 
        println("usage : scala Pouring [filename]. you can use input.txt as an input. Or you can just run 'scala Pouring' and enter custom values")
    
    
    def main (args: Array[String]) = {

        if (args.length == 1) {
            
            val filename = args(0)
            val f = new File(filename)
            if (!f.exists()) {
                println(filename + "doesn't exists")
                System.exit(-1)
            }
            
            val cur = System.currentTimeMillis 
        
            println("Running all problems from " + filename)
            
            Source.fromFile(filename)
                  .getLines
                  .map(line_to_int(_))
                  .grouped(4)
                  .filter(isProblem(_))
                  .foreach(x=> { println(); solve(x) })
            
            println("Runned all the problems from " + filename + " in " + (System.currentTimeMillis - cur) + "ms")
            System.exit(0)
        }
        
        else if(args.length == 0) {
            solve(Array("how many containers ?", 
                        "capacities ? (space separated)", 
                        "starting values ?", 
                        "goal ? (zeros will be discarded)")
                   .map(q => line_to_int(askForInput(q))))
            System.exit(0)
        }
        
        else {
            usage()
            System.exit(-1)
        }
    }
}