How do I write a ZipN-like function in F#?

Question

I want to create a function with the signature seq<#seq<'a>> ->seq<seq<'a>> that acts like a Zip method taking a sequence of an arbitrary number of input sequences (instead of 2 or 3 as in Zip2 and Zip3) and returning a sequence of sequences instead of tuples as a result.

That is, given the following input:

[[1;2;3];
 [4;5;6];
 [7;8;9]]

it will return the result: [[1;4;7]; [2;5;8]; [3;6;9]]

except with sequences instead of lists.

I am very new to F#, but I have created a function that does what I want, but I know it can be improved. It's not tail recursive and it seems like it could be simpler, but I don't know how yet. I also haven't found a good way to get the signature the way I want (accepting, e.g., an int list list as input) without a second function.

I know this could be implemented using enumerators directly, but I'm interested in doing it in a functional manner.

Here's my code:

let private Tail seq = Seq.skip 1 seq
let private HasLengthNoMoreThan n = Seq.skip n >> Seq.isEmpty

let rec ZipN_core = function
    | seqs when seqs |> Seq.isEmpty -> Seq.empty
    | seqs when seqs |> Seq.exists Seq.isEmpty -> Seq.empty
    | seqs ->
        let head = seqs |> Seq.map Seq.head
        let tail = seqs |> Seq.map Tail |> ZipN_core
        Seq.append (Seq.singleton head) tail

// Required to change the signature of the parameter from seq<seq<'a> to seq<#seq<'a>>
let ZipN seqs = seqs |> Seq.map (fun x -> x |> Seq.map (fun y -> y)) |> ZipN_core

Answer 1

let zipn items = items |> Matrix.Generic.ofSeq |> Matrix.Generic.transpose

Or, if you really want to write it yourself:

let zipn items = 
  let rec loop items =
    seq {
      match items with
      | [] -> ()
      | _ -> 
        match zipOne ([], []) items with
        | Some(xs, rest) -> 
          yield xs
          yield! loop rest
        | None -> ()
    }
  and zipOne (acc, rest) = function
    | [] -> Some(List.rev acc, List.rev rest)
    | []::_ -> None
    | (x::xs)::ys -> zipOne (x::acc, xs::rest) ys
  loop items

Answer 2

Since this seems to be the canonical answer for writing a zipn in f#, I wanted to add a "pure" seq solution that preserves laziness and doesn't force us to load our full source sequences in memory at once like the Matrix.transpose function. There are scenarios where this is very important because it's a) faster and b) works with sequences that contain 100s of MBs of data!

This is probably the most un-idiomatic f# code I've written in a while but it gets the job done (and hey, why would there be sequence expressions in f# if you couldn't use them for writing procedural code in a functional language).

 let seqdata = seq {
  yield Seq.ofList [ 1; 2; 3 ]
  yield Seq.ofList [ 4; 5; 6 ]
  yield Seq.ofList [ 7; 8; 9 ]
}

let zipnSeq (src:seq<seq<'a>>) = seq {
  let enumerators = src |> Seq.map (fun x -> x.GetEnumerator()) |> Seq.toArray
  if (enumerators.Length > 0) then
    try 
      while(enumerators |> Array.forall(fun x -> x.MoveNext())) do 
        yield enumerators |> Array.map( fun x -> x.Current)
    finally 
      enumerators |> Array.iter (fun x -> x.Dispose())
}

zipnSeq seqdata |> Seq.toArray


val it : int [] [] = [|[|1; 4; 7|]; [|2; 5; 8|]; [|3; 6; 9|]|]

By the way, the traditional matrix transpose is much more terse than @Daniel's answer. Though, it requires a list or LazyList that both will eventually have the full sequence in memory.

let rec transpose = 
  function 
  | (_ :: _) :: _ as M -> List.map List.head M :: transpose (List.map List.tail M)
  | _ -> []

Answer 3

To handle having sub-lists of different lengths, I've used option types to spot if we've run out of elements.

let split = function
    | []    -> None,    []
    | h::t  -> Some(h), t

let rec zipN listOfLists =
    seq { let splitted = listOfLists |> List.map split

          let anyMore = splitted |> Seq.exists (fun (f, _) -> f.IsSome)

          if anyMore then
              yield splitted |> List.map fst
              let rest = splitted |> List.map snd
              yield! rest |> zipN }

This would map

let ll = [ [ 1; 2; 3 ];
           [ 4; 5; 6 ];
           [ 7; 8; 9 ] ]

to

seq
    [seq [Some 1; Some 4; Some 7]; seq [Some 2; Some 5; Some 8];
     seq [Some 3; Some 6; Some 9]]

and

let ll = [ [ 1; 2; 3 ];
           [ 4; 5; 6 ];
           [ 7; 8 ] ]

to

seq
    [seq [Some 1; Some 4; Some 7]; seq [Some 2; Some 5; Some 8];
     seq [Some 3; Some 6; null]]

This takes a different approach to yours, but avoids using some of the operations that you had before (e.g. Seq.skip, Seq.append), which you should be careful with.