CODEWITHSUNDEEP
scala
learningTheRopes
learningTheRopes

Reputation: 75

Scala : How to break out of a nested for comprehension

I'm trying to write some code as below - 

def kthSmallest(matrix: Array[Array[Int]], k: Int): Int = {
  val pq = new PriorityQueue[Int]() //natural ordering
  var count = 0

  for (
    i <- matrix.indices;
    j <- matrix(0).indices
  ) yield {
      pq += matrix(i)(j)
      count += 1
  } //This would yield Any!

  pq.dequeue() //kth smallest. 

}

My question is, that I only want to loop till the time count is less than k (something like takeWhile(count != k)), but as I'm also inserting elements into the priority queue in the yield, this won't work in the current state.

My other options are to write a nested loop and return once count reaches k. Is it possible to do with yield? I could not find a idiomatic way of doing it yet. Any pointers would be helpful.

Upvotes: 0

Views: 1068

Answers (2)

Sergey
Sergey

Reputation: 2900

It's not idiomatic for Scala to use vars or break loops. You can go for recursion, lazy evaluation or duct tape a break, giving up on some performance (just like return, it's implemented as an Exception, and won't perform well enough). Here are the options broken down:

  1. Use recursion - recursive algorithms are the analog of loops in functional languages

    def kthSmallest(matrix: Array[Array[Int]], k: Int): Int = {
  val pq = new PriorityQueue[Int]() //natural ordering

  @tailrec
  def fillQueue(i: Int, j: Int, count: Int): Unit =
    if (count >= k || i >= matrix.length) ()
    else {
      pq += matrix(i)(j)
      fillQueue(
        if (j >= matrix(i).length - 1) i + 1 else i,
        if (j >= matrix(i).length - 1) 0 else j + 1,
        count + 1)
    }

  fillQueue(0, 0, 0)

  pq.dequeue() //kth smallest.
}

  2. Use a lazy structure, as chengpohi suggested - this doesn't sound very much like a pure function though. I'd suggest to use an Iterator instead of a Stream in this case though - as iterators don't memoize the steps they've gone through (might spare some memory for large matrices).

  3. For those desperately willing to use break, Scala supports it in an attachable fashion (note the performance caveat mentioned above):

    import scala.util.control.Breaks

breakable {
  // loop code
  break
}

Upvotes: 1

chengpohi
chengpohi

Reputation: 14217

There is a way using the Stream lazy evaluation to do this. Since for yield is equal to flatMap, you can convert for yield to flatMap with Stream:

matrix.indices.toStream.flatMap(i => {
  matrix(0).indices.toStream.map(j => {
    pq += matrix(i)(j)
    count += 1
  })
}).takeWhile(_ => count <= k)

Use toStream to convert the collection to Stream, and Since Stream is lazy evaluation, so we can use takeWhile to predicate count to terminate the less loops without init others.

Upvotes: 1

Related Questions