CODEWITHSUNDEEP
scalacollections
Piotr Kopeć
Piotr Kopeć

Reputation: 13

Break flatMap after first Some occurence

abstract class Node(id: Long, name: String) {
  def find(id: Long) = if (this.id == id) Some(this) else None
}
case class Contact(id: Long, name: String, phone: String) extends Node(id, name)
case class Group(id: Long, name: String, entries: Node*) extends Node(id, name) {
  override def find(id: Long): Option[Node] = {
    super.find(id) orElse entries.flatMap(e => e.find(id)).headOption
  }
}
val tree =
  Group(0, "Root"
      , Group(10, "Shop", Contact(11, "Alice", "312"))
      , Group(20, "Workshop"
          , Group(30, "Tyres", Contact(31, "Bob", "315"), Contact(32, "Greg", "319"))
          , Contact(33, "Mary", "302"))
      , Contact(1, "John", "317"))
println(tree.find(32))

Tree data is built from Contacts and Groups (w/ sub-Groups and Contacts). I want to find a node with specific id. Currently I traverse Group's members using:

entries.flatMap(e => e.find(id)).headOption

but it isn't optimal because I check all child entries instead of break upon first finding.

I'd appreciate your help in magic Scala Wold. Thanks.

Upvotes: 1

Views: 890

Answers (2)

Marek Adamek
Marek Adamek

Reputation: 520

Another way to approach this problem could be providing traverse support for the data structure. Overall it's a tree so it can be traversed easily. Please check the code below:

sealed abstract class Node(val id: Long, val name: String) extends Traversable[Node] {
  def foreach[U](f:Node => U) = this match {
    case x: Contact => f(x)
    case x: Group => f(x); x.entries.foreach(_ foreach f)
  }
}
case class Contact(override val id: Long, override val name: String, phone: String) extends Node(id, name) {
  override def toString = (id, name, phone).toString
}
case class Group(override val id: Long, override val name: String, entries: Node*) extends Node(id, name) {
  override def toString = (id, name).toString + entries.map(_.toString).mkString
}

val tree =
  Group(0, "Root"
    , Group(10, "Shop", Contact(11, "Alice", "312"))
    , Group(20, "Workshop"
      , Group(30, "Tyres", Contact(31, "Bob", "315"), Contact(32, "Greg", "319"))
      , Contact(33, "Mary", "302"))
    , Contact(1, "John", "317"))

println(tree) // (0,Root)(10,Shop)(11,Alice,312)(20,Workshop)(30,Tyres)(31,Bob,315)(32,Greg,319)(33,Mary,302)(1,John,317)
println(tree find {_.id == 32}) //Some((32,Greg,319))
println(tree map {_.name }) //List(Root, Shop, Alice, Workshop, Tyres, Bob, Greg, Mary, John)

The good thing is that now you can use of all the benefits of Traversable[Node] trait. The problem on the other hand is that I had to override toString method in case classes. So I guess there is still room for improvement.

Upvotes: 0

Michael Zajac
Michael Zajac

Reputation: 55569

You want collectFirst, which will select the first matching element and wrap it in Some, or None if it's not found. You can also turn entries into a view, to make the evaluation lazy.

entries.view.map(_.find(id)).collectFirst { case Some(node) => node }

It will work with your original code, as well:

entries.view.flatMap(_.find(id)).headOption

Upvotes: 3

Related Questions