Swift 4
The general sequence case
Citing the OP:
Even more general would be to have a function that takes any sequence
and provides a sequence that's the running total of the input
sequence.
Consider some arbitrary sequence (conforming to Sequence
), say
var seq = 1... // 1, 2, 3, ... (CountablePartialRangeFrom)
To create another sequence which is the (lazy) running sum over seq
, you can make use of the global sequence(state:next:)
function:
var runningSumSequence =
sequence(state: (sum: 0, it: seq.makeIterator())) { state -> Int? in
if let val = state.it.next() {
defer { state.sum += val }
return val + state.sum
}
else { return nil }
}
// Consume and print accumulated values less than 100
while let accumulatedSum = runningSumSequence.next(),
accumulatedSum < 100 { print(accumulatedSum) }
// 1 3 6 10 15 21 28 36 45 55 66 78 91
// Consume and print next
print(runningSumSequence.next() ?? -1) // 120
// ...
If we'd like (for the joy of it), we could condense the closure to sequence(state:next:)
above somewhat:
var runningSumSequence =
sequence(state: (sum: 0, it: seq.makeIterator())) {
(state: inout (sum: Int, it: AnyIterator<Int>)) -> Int? in
state.it.next().map { (state.sum + $0, state.sum += $0).0 }
}
However, type inference tends to break (still some open bugs, perhaps?) for these single-line returns of sequence(state:next:)
, forcing us to explicitly specify the type of state
, hence the gritty ... in
in the closure.
Alternatively: custom sequence accumulator
protocol Accumulatable {
static func +(lhs: Self, rhs: Self) -> Self
}
extension Int : Accumulatable {}
struct AccumulateSequence<T: Sequence>: Sequence, IteratorProtocol
where T.Element: Accumulatable {
var iterator: T.Iterator
var accumulatedValue: T.Element?
init(_ sequence: T) {
self.iterator = sequence.makeIterator()
}
mutating func next() -> T.Element? {
if let val = iterator.next() {
if accumulatedValue == nil {
accumulatedValue = val
}
else { defer { accumulatedValue = accumulatedValue! + val } }
return accumulatedValue
}
return nil
}
}
var accumulator = AccumulateSequence(1...)
// Consume and print accumulated values less than 100
while let accumulatedSum = accumulator.next(),
accumulatedSum < 100 { print(accumulatedSum) }
// 1 3 6 10 15 21 28 36 45 55 66 78 91
The specific array case: using reduce(into:_:)
As of Swift 4, we can use reduce(into:_:)
to accumulate the running sum into an array.
let runningSum = arr
.reduce(into: []) { $0.append(($0.last ?? 0) + $1) }
// [2, 4, 6, 8, 10, 12]
By using reduce(into:_:)
, the [Int]
accumulator will not be copied in subsequent reduce iterations; citing the Language reference:
This method is preferred over reduce(_:_:)
for efficiency when the
result is a copy-on-write type, for example an Array
or a
Dictionary
.
See also the implementation of reduce(into:_:)
, noting that the accumulator is provided as an inout
parameter to the supplied closure.
However, each iteration will still result in an append(_:)
call on the accumulator array; amortized O(1)
averaged over many invocations, but still an arguably unnecessary overhead here as we know the final size of the accumulator.
Because arrays increase their allocated capacity using an exponential
strategy, appending a single element to an array is an O(1)
operation
when averaged over many calls to the append(_:)
method. When an array
has additional capacity and is not sharing its storage with another
instance, appending an element is O(1)
. When an array needs to
reallocate storage before appending or its storage is shared with
another copy, appending is O(n)
, where n is the length of the array.
Thus, knowing the final size of the accumulator, we could explicitly reserve such a capacity for it using reserveCapacity(_:)
(as is done e.g. for the native implementation of map(_:)
)
let runningSum = arr
.reduce(into: [Int]()) { (sums, element) in
if let sum = sums.last {
sums.append(sum + element)
}
else {
sums.reserveCapacity(arr.count)
sums.append(element)
}
} // [2, 4, 6, 8, 10, 12]
For the joy of it, condensed:
let runningSum = arr
.reduce(into: []) {
$0.append(($0.last ?? ($0.reserveCapacity(arr.count), 0).1) + $1)
} // [2, 4, 6, 8, 10, 12]
Swift 3: Using enumerated()
for subsequent calls to reduce
Another Swift 3 alternative (with an overhead ...) is using enumerated().map
in combination with reduce
within each element mapping:
func runningSum(_ arr: [Int]) -> [Int] {
return arr.enumerated().map { arr.prefix($0).reduce($1, +) }
} /* thanks @Hamish for improvement! */
let arr = [2, 2, 2, 2, 2, 2]
print(runningSum(arr)) // [2, 4, 6, 8, 10, 12]
The upside is you wont have to use an array as the collector in a single reduce
(instead repeatedly calling reduce
).