How to fix my implementation of partial sums in Scala?

Question

This is a follow-up to my previous question. I would like to implement s.scanLeft(0)(_ + _) by myself (just as an exercise)

That is, I would like to write function partial_sums, which receives stream s = s1, s2, s3, ... and produces a new stream s1, s1 + s2, s1 + s2 + s3, ...

I implement it as follows:

def add_streams(s1:Stream[Int], s2:Stream[Int]) =
  (s1 zip s2) map {case (x, y) => x + y}

def partial_sums(s:Stream[Int]):Stream[Int] =
  Stream.cons(s.head, add_streams(partial_sums(s), s.tail))

This code works ok. However it looks like it takes O(n) to get the n-th element of partial_sums. (i.e. s[1] + s[2] + s[3] ... + s[n]). I would like to code partial_sums[n] = partial_sums[n-1] + s[n], which takes O(1) to calculate the n-th element.

Is it correct? How would you fix the code?

Answer 1

The basic idea is to keep a running total, rather than adding streams in bulk

def partialSums(s:Stream[Int]) = if(s.isEmpty) new Stream[Int]() else partialSums(s, 0)

def partialSums(s:Stream[Int], runningTotal:Int)= Stream.cons(s.head+runningTotal, partialSums(s.tail, s.head+runningTotal)