Irregular segmented scan
Accepts an irregular array encoded as a value vector and a flag vector, then performs a segmented scan. A true flag denotes the beginning of a new segment.
def segscan 't [n]
(op: t -> t -> t)
(ne: t)
(flags: [n]bool)
(vals: [n]t) : [n]t =
let pairs =
scan (\(v1, f1) (v2, f2) ->
let f = f1 || f2
let v = if f2 then v2 else op v1 v2
in (v, f))
(ne, false)
(zip vals flags)
let (res, _) = unzip pairs
in res
def segprefix_sum = segscan (+) 0i32> segprefix_sum [true, true, false, true, false, false]
[1,2,3,4,5,6][1i32, 2i32, 5i32, 4i32, 9i32, 15i32]