5.2

Write the function take(n) for returning the first n elements of a Stream, and drop(n) for skipping the first n elements of a Stream.

Note: The key to remember here is that since in take we need to create another stream and keep the elements in order. While in drop, we drop from the head of the stream and return the rest

Solution

object Solution extends App {
  val stream: Stream[Int] = Stream(1,2,3,4,5,6)
  println("Stream:" + stream.toList)

  val take3: Stream[Int] = stream.take(3)
  println("take(3):" + take3.toList) 

  val drop3: Stream[Int] = stream.drop(3)
  println("drop(3):" + drop3.toList) 
}

import scala.annotation
trait Stream[+A] {
  import Stream._
  def take(n: Int): Stream[A] = this match {
    case Empty => empty
    case Cons(h, t) if n > 1 => cons(h(), t().take(n-1))
    case Cons(h, t) if n == 1 => cons(h(), empty)
  }

  def drop(n: Int): Stream[A] = {

    @annotation.tailrec
    def go(s: Stream[A], n: Int): Stream[A] = s match {
      case Empty => s
      case Cons(h, t) if n >= 1 => go(t(), n - 1)
      case Cons(h, t) if n == 0 => s
    }

    go(this, n)
  }

  def toList: List[A] = {
    @annotation.tailrec
    def go(l: List[A], s: Stream[A]): List[A] = s match {
      case Empty => l.reverse
      case Cons(h,t) => go(h()::l, t())
    }

    go(List(), this)
  }
}


case object Empty extends Stream[Nothing]
case class Cons[+A](h: () => A, t: () => Stream[A]) extends Stream[A]
object Stream {
  def empty[A]: Stream[A] = Empty

  def cons[A](hd: => A, tl: => Stream[A]): Stream[A] = {
    lazy val head = hd
    lazy val tail = tl
    Cons(() => head, () => tail)
  }

  def apply[A](as: A*): Stream[A] = 
    if(as.isEmpty) empty else cons(as.head, apply(as.tail: _*))
}

Output

Stream:List(1, 2, 3, 4, 5, 6) 
take(3):List(1, 2, 3) 
drop(3):List(4, 5, 6)