CODEWITHSUNDEEP
swift
mistercake
mistercake

Reputation: 723

Elegant way to get the first n elements of a SequenceType

Is there an elegant way (think one-liner) to get the first n elements of a SequenceType in Swift?

I could of course write a for-loop that terminates after n elements but that's a little bulky.

Note also that the solution should be able to deal with infinite sequences.

Upvotes: 5

Views: 5396

Answers (4)

DeFrenZ
DeFrenZ

Reputation: 2252

Isn't it exactly what mySequence.prefix(numberOfElements) does?

Upvotes: 21

Iain Bryson
Iain Bryson

Reputation: 410

Why not

var seq = NominalSequence().generate()
var a = (0..<10).map({_ in seq.next()!})

?

Two lines, but funtional-ish.

Upvotes: -1

Rob Napier
Rob Napier

Reputation: 299345

In Swift 2, you can create this as an extension:

extension SequenceType {
    func take(n: Int) -> [Generator.Element] {
        var result: [Generator.Element] = []
        var g = self.generate()
        for _ in 1...n {
            if let next = g.next() {
                result.append(next)
            } else {
                break
            }
        }
        return result
    }
}

In Swift 1, it would have to written as a function:

func take<Seq: SequenceType>(n: Int, xs: Seq) -> [Seq.Generator.Element] {
    var result: [Seq.Generator.Element] = []
    var g = xs.generate()
    for _ in 1...n {
        if let next = g.next() {
            result.append(next)
        } else {
            break
        }
    }
    return result
}

Note that in either case, SequenceType does not specify what happens if you call generate() more than once. It could return the same values (as in an Array). It could return different values (as in an audio data stream). It could return nothing at all. So the caller of take() may need some special knowledge about its impact on the sequence.

Upvotes: 2

GoZoner
GoZoner

Reputation: 70175

SequenceType only supports generate(). Perhaps a more 'Swiftly' approach would be to define a Generator that given a start and end index and a 'base' generator would skip over the first elements, start returning some, and then stop after end index. Like this:

struct SubscriptGenerator<Base: GeneratorType> : GeneratorType {
  var nowIndex : Int = 0
  let endIndex : Int
  let begIndex : Int
  var generator : Base

  init (generator: Base, startIndex: Int, endIndex: Int) {
    precondition(startIndex < endIndex, "oops")
    self.generator = generator
    self.endIndex  = endIndex
    self.begIndex  = startIndex
  }

  // MARK - GeneratorType

  typealias Element = Base.Element

  mutating func next() -> Element? {
    while (nowIndex < begIndex) { nowIndex++; generator.next () }
    return nowIndex++ < endIndex ? generator.next() : nil
  }
}

This is only an example. One could define an convenience init() that takes a SequenceType and produces the base generator. In action:

75> var gen = [10,20,30,40,50].generate()
76> var sg = SubscriptGenerator(generator: gen, startIndex: 1, endIndex:3) 
sg: SubscriptGenerator<IndexingGenerator<[Int]>> = { ... }
 77> sg.next()
$R2: Int? = 20
 78> sg.next()
$R3: Int? = 30
 79> sg.next()
$R4: Int? = nil

See Swift's EnumerateGenerator for an example.

Note: it might be the Stride nexus of Swift functionality does what you want already.

Upvotes: 0

Related Questions