38a1fbc6688e

Enhance `prefix-sums`
[view raw] [browse files]
author Steve Losh <steve@stevelosh.com>
date Mon, 09 Nov 2020 21:51:32 -0400
parents 1d3928ec556d
children 2db6543352b6 f56991dfcaea
branches/tags (none)
files src/sequences.lisp

Changes

--- a/src/sequences.lisp	Sat Oct 24 15:55:36 2020 -0400
+++ b/src/sequences.lisp	Mon Nov 09 21:51:32 2020 -0400
@@ -9,19 +9,24 @@
   `(integer 0 (,length)))
 
 
-(defun prefix-sums (sequence)
-  "Return a list of the prefix sums of the numbers in `sequence`.
+(defun prefix-sums (sequence &key key (result-type 'list))
+  "Return the prefix sums of the elements of `sequence`.
+
+  If `key` is given, it will be called on the elements before summing.
+  `result-type` must be a type suitable for passing to `map`.
 
   Example:
 
     (prefix-sums '(10 10 10 0 1))
-    => (10 20 30 30 31)
+    ; => (10 20 30 30 31)
+
+    (prefix-sums \"ABCD\" :key #'char-code :result-type '(vector fixnum))
+    ; => #(65 131 198 266)
 
   "
-  (iterate
-    (for i :in-whatever sequence)
-    (sum i :into s)
-    (collect s)))
+  (let ((sum 0))
+    (map result-type (lambda (x) (incf sum (if key (funcall key x) x)))
+         sequence)))
 
 (defun frequencies (sequence &key (test #'eql) key)
   "Return a hash table containing the frequencies of the elements of `sequence`.