Search code examples
f#sequences

How to split a sequence in F# based on another sequence in an idiomatic way


I have, in F#, 2 sequences, each containing distinct integers, strictly in ascending order: listMaxes and numbers.

If not Seq.isEmpty numbers, then it is guaranteed that not Seq.isEmpty listMaxes and Seq.last listMaxes >= Seq.last numbers.

I would like to implement in F# a function that returns a list of list of integers, whose List.length equals Seq.length listMaxes, containing the elements of numbers divided in lists, where the elements of listMaxes limit each group.

For example: called with the arguments

listMaxes = seq [ 25; 56; 65; 75; 88 ]
numbers = seq [ 10; 11; 13; 16; 20; 25; 31; 38; 46; 55; 65; 76; 88 ]

this function should return

[ [10; 11; 13; 16; 20; 25]; [31; 38; 46; 55]; [65]; List.empty; [76; 88] ]

I can implement this function, iterating over numbers only once:

let groupByListMaxes listMaxes numbers = 
    if Seq.isEmpty numbers then
        List.replicate (Seq.length listMaxes) List.empty
    else
        List.ofSeq (seq {
            use nbe = numbers.GetEnumerator ()
            ignore (nbe.MoveNext ())
            for lmax in listMaxes do
                yield List.ofSeq (seq {
                    if nbe.Current <= lmax then
                        yield nbe.Current
                    while nbe.MoveNext () && nbe.Current <= lmax do
                        yield nbe.Current
                })
        })

But this code feels unclean, ugly, imperative, and very un-F#-y.

Is there any functional / F#-idiomatic way to achieve this?


Solution

  • It can be done like this with pattern matching and tail recursion:

    let groupByListMaxes listMaxes numbers =
        let rec inner acc numbers =
            function
            | [] -> acc |> List.rev
            | max::tail ->
                let taken = numbers |> Seq.takeWhile ((>=) max) |> List.ofSeq
                let n = taken |> List.length
                inner (taken::acc) (numbers |> Seq.skip n) tail
    
        inner [] numbers (listMaxes |> List.ofSeq)
    

    Update: I also got inspired by fold and came up with the following solution that strictly refrains from converting the input sequences.

    let groupByListMaxes maxes numbers =
        let rec inner (acc, (cur, numbers)) max = 
            match numbers |> Seq.tryHead with
            // Add n to the current list of n's less
            // than the local max
            | Some n when n <= max ->
                let remaining = numbers |> Seq.tail  
                inner (acc, (n::cur, remaining)) max
            // Complete the current list by adding it
            // to the accumulated result and prepare
            // the next list for fold.
            | _ ->
                (List.rev cur)::acc, ([], numbers)
    
        maxes |> Seq.fold inner ([], ([], numbers)) |> fst |> List.rev