I wrote this solution to Kattis problem Amanda Lounges in Scala. TL;DR: It's a graph theory problem where I read in a list of edges from stdin and try to compute the minimum number of nodes that should be "lounges," given that each edge has a weight of 0, 1, or 2, and an edge with a weight of n
must have n
of its endpoints as lounges. The general idea is that, once you pick a node in a component, it can either have a lounge or not, and, from there, everything in that component is determined. I'm relatively new to Scala, so I'm looking for general feedback as well as how I can use less memory (right now, when I submit it to Kattis, it exceeds the 1024 MB memory limit for the problem). I'm trying to do it functionally, though I do make one compromise for speed.
import scala.collection.mutable.{ Map => MMap }
import scala.annotation.tailrec
import scala.io.Source
object Main extends App {
val lines = Source.stdin.getLines()
val Array(n, m) = lines.next().split("\\s+").map(_.toInt)
// groupBy was too slow
val graph = MMap[Int, List[Edge]]()
for (e <- lines.flatMap(parseEdge)) {
if (!(graph contains e.start)) {
graph(e.start) = Nil
}
graph(e.start) = e :: graph(e.start)
}
println(numLoungesPossible().getOrElse("impossible"))
@tailrec
def numLoungesPossible(curNode: Int = 1, numSoFar: Int = 0,
vis: Set[Int] = Set()): Option[Int] = {
if (curNode > n) {
Some(numSoFar)
} else if (vis contains curNode) {
numLoungesPossible(curNode + 1, numSoFar, vis)
} else {
val num = numLoungesInComp(curNode)
val newVis = component(curNode, vis)
num match {
case None => None
case Some(i) => numLoungesPossible(curNode + 1, numSoFar + i, newVis)
}
}
}
/**
* Builds a set of the nodes in the graph that are connected to a given node.
* Can operate by adding them to an existing set if that optional parameter is
* provided.
*/
def component(node: Int, vis: Set[Int] = Set()): Set[Int] = {
if (vis contains node)
vis
else {
val newVis = vis + node
graph.getOrElse(node, Nil).foldLeft(newVis) { (curVis, n) =>
component(n.end, curVis)
}
}
}
/**
* Computes the minimum number of lounges in the component of the graph
* connected to the given node.
*/
def numLoungesInComp(node: Int): Option[Int] = {
val n1 = numLounges(node, true)
val n2 = numLounges(node, false)
val ns = n1.toList ::: n2.toList
ns map (_._1) reduceOption math.min
}
/**
* Computes the number of lounges in the component of the graph connected to
* the given node if a lounge is placed here or if a lounge is not placed
* here, depending on whether the parameter `lounge` is true or false,
* respectively.
*/
def numLounges(node: Int, lounge: Boolean,
vis: Map[Int, Boolean] = Map()): Option[(Int, Map[Int, Boolean])] = {
val neighbors = graph.getOrElse(node, Nil)
if (vis contains node) {
// This is visited; check for conflict
if (vis(node) == lounge) Some((0, vis)) else None
} else if (neighbors.exists(_.lounges == (if (lounge) 0 else 2))) {
None
} else {
lazy val nLoungesHere = if (lounge) 1 else 0
val ret = computeLounges(neighbors, lounge, vis + (node -> lounge))
ret map { case (n, v) => (n + nLoungesHere, v) }
}
}
/**
* Compute the number of lounges in the given neighbors of a node and all
* other unvisited nodes connected to them if a lounge is or is not placed on
* the root node (depending on the value of `lounge`).
*/
@tailrec
def computeLounges(neighbors: List[Edge], lounge: Boolean,
vis: Map[Int, Boolean], numSoFar: Int = 0):
Option[(Int, Map[Int, Boolean])] = neighbors match {
case Nil => Some((numSoFar, vis))
case e :: rest => {
// Whether or not to place a lounge at this neighbor
val nLounge = if (e.lounges == 1) !lounge else lounge
numLounges(e.end, nLounge, vis) match {
case None => None
case Some((i, newVis)) =>
computeLounges(rest, lounge, newVis, numSoFar + i)
}
}
}
def parseEdge(line: String) = {
val Array(s, e, n) = line.split("\\s+").map(_.toInt)
val edge = Edge(s, e, n)
List(edge, edge.reversed)
}
}
case class Edge(start: Int, end: Int, lounges: Int) {
def reversed = copy(start = end, end = start)
}