0

This is piggybacking on the classic question how to do Seq.takeWhile + one item in F# by extending it to the general case of extracting n items after the predicate returns false.

The obvious bare-bone solution may be based on Tomas' answer to the original question by equipping the recursive loop on IEnumerator<'T> with an additional counter, but I can think of at least two other approaches.

First there's Seq.pairwise suggested by Be Brave Be Like Ukraine which could be replicated as often as needed; the overhead of the additional tuples will probably mean that it's not very efficient:

// Trivial case for n = 0
let takeWhile pred =
    Seq.map (fun t -> pred t, Some t)
    >> Seq.takeWhile fst
    >> Seq.choose snd

let takeWhile1 pred =
    Seq.map (fun t -> pred t, Some t)
    >> Seq.append (Seq.singleton (true, None))
    >> Seq.pairwise
    >> Seq.takeWhile (fst >> fst)
    >> Seq.choose (snd >> snd)

let takeWhile2 pred =
    Seq.map (fun t -> pred t, Some t)
    >> Seq.append (Seq.singleton (true, None))
    >> Seq.append (Seq.singleton (true, None))
    >> Seq.pairwise
    >> Seq.pairwise
    >> Seq.takeWhile (fst >> fst >> fst)
    >> Seq.choose (snd >> snd >> snd)

[1..7] |> takeWhile ((>=) 3) |> Seq.toList
// val it : int list = [1; 2; 3]
[1..7] |> takeWhile1 ((>=) 3) |> Seq.toList
// val it : int list = [1; 2; 3; 4]
[1..7] |> takeWhile2 ((>=) 3) |> Seq.toList
// val it : int list = [1; 2; 3; 4; 5]

The other involves deconstruction of the sequence by Seq.tryHead and Seq.tail. After testing I realize that this would be sub-optimal as the number of times the elements of the sequence get evaluated grows quadratically.

let takeWhileN pred n source = 
    let rec aux i source = seq{
        match Seq.tryHead source with
        | Some x when pred x && i = 0 ->
            yield x
            yield! aux 0 (Seq.tail source)
        | Some x when i < n ->
            yield x
            yield! aux (i + 1) (Seq.tail source)
        | _ -> () }
    aux 0 source

[1..7] |> takeWhileN ((>=) 3) 0 |> Seq.toList
// val it : int list = [1; 2; 3]
[1..7] |> takeWhileN ((>=) 3) 1 |> Seq.toList
// val it : int list = [1; 2; 3; 4]
[1..7] |> takeWhileN ((>=) 3) 2 |> Seq.toList
// val it : int list = [1; 2; 3; 4; 5]

Can this be improved upon?

kaefer
  • 5,491
  • 1
  • 15
  • 20

2 Answers2

2

Your solution is very good. The only note is that once you're past the predicate, you can use Seq.truncate rather than yielding items one at a time.

let rec takeWhileN pred n source = seq {
    match Seq.tryHead source with
    | Some x when pred x ->
        yield x
        yield! takeWhileN pred n (Seq.tail source)
    | _ -> yield! Seq.truncate n source            
}
Sanity Challenged
  • 495
  • 3
  • 6
0

You can also just be direct about it:

let takeWhileN pred n source =
    let matches = source |> Seq.takeWhile pred |> Seq.length
    source |> Seq.truncate (matches + n)
Sanity Challenged
  • 495
  • 3
  • 6