import scala.io.Source

def findGroup(v: String, groups: Map[String, Set[String]]): Set[String] = {
  def loop(toVisit: Set[String], visited: Set[String]): Set[String] = (toVisit -- visited) match {
    case v if v.isEmpty => visited
    case v => loop(v.flatMap(n => groups.get(n).getOrElse(Set.empty[String])), visited ++ v)
  }
  loop(Set(v), Set.empty[String])
}

def partOne(groups: Map[String, Set[String]]): Int = findGroup("0", groups).size
def partTwo(groups: Map[String, Set[String]]): Int = groups.foldLeft((0, Set.empty[String])){ case ((n, visited), (k, _)) =>
  if (!visited.contains(k)) (n + 1, visited ++ findGroup(k, groups))
  else (n, visited)
}._1

val input = Source.fromFile("../input12.txt").mkString.split("\n").toSeq.map { l =>
  l.split(" <-> ").toList match {
    case p :: c :: Nil => (p -> c.split(", ").toSet)
    case _ => throw new RuntimeException("bad input")
  }
}.toMap

println(s"Part 1: ${partOne(input)}")
println(s"Part 2: ${partTwo(input)}")