6
\$\begingroup\$

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)
}
\$\endgroup\$
  • \$\begingroup\$ I don't know Scala so I won't be able to help on the code quality, but I do see some potential improvements in the algorithm. The first thing you should do is to preprocess the graph by picking (or non-picking) all the nodes incident to edges of weight 2 or 0, and then verify that the remaining graph is bipartite. Then, for each connected component, you take the minimal side as lounges, i.e., color each of the vertices "red" or "blue" depending on side in bipartition and pick the smallest side. \$\endgroup\$ – Pål GD Apr 18 '17 at 18:49
  • \$\begingroup\$ @PålGD would that be significantly faster or use significantly less memory? \$\endgroup\$ – Brian McCutchon Apr 19 '17 at 3:29
  • \$\begingroup\$ Yes, this problem should be solvable in linear time, and it looks as if this algorithm might do exponentially much work, no? \$\endgroup\$ – Pål GD Apr 19 '17 at 15:53
  • \$\begingroup\$ @PålGD I think not. I do two linear passes over each connected component, and I don't look at a component twice. Unless I'm missing something... \$\endgroup\$ – Brian McCutchon Apr 19 '17 at 19:05

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.