Swift3で配列のシャッフルをLazyに行って計算量を減らす

全てシャッフルする場合

Swift3で配列のシャッフルを行うコードを検索したところswift3 - Shuffle array swift 3 - Stack Overflowがヒットしました。

このコードを使ってシャッフルしたのち最初の2個の値を取得するコードは以下のようになります。同時に実行速度を測ってみます。

extension MutableCollection where Index == Int {
    /// Shuffle the elements of `self` in-place.
    mutating func shuffle() {
        // empty and single-element collections don't shuffle
        if count < 2 { return }

        for i in startIndex ..< endIndex - 1 {
            let j = Int(arc4random_uniform(UInt32(endIndex - i))) + i
            if i != j {
                swap(&self[i], &self[j])
            }
        }
    }
}

extension Collection {
    /// Return a copy of `self` with its elements shuffled
    func shuffled() -> [Iterator.Element] {
        var list = Array(self)
        list.shuffle()
        return list
    }
}

let array = Array(0...1000000)

printTimeElapsedWhenRunningCode(title: "Shuffle all") {
    _ = array.shuffled()[0...1]
}

結果はTime elapsed for Shuffle all: 79.2173529863358 sと出ました。かなり時間がかかっています。

Lazyにシャッフルする場合

これをLazyに行って計算量を減らします。まずShuffledIteratorを作成します。次にLazySequenceProtocolに適合するLazyShuffledSequenceを作成します。

struct ShuffledIterator<T>: IteratorProtocol {

    typealias Element = T

    private var elements: [T]

    init(elements: [T]) {
        self.elements = elements
    }

    mutating func next() -> T? {
        guard elements.count > 0 else { return nil }
        let index = arc4random_uniform(UInt32(elements.count))
        return elements.remove(at: Int(index))
    }

}

struct LazyShuffledSequence<T>: LazySequenceProtocol {

    func makeIterator() -> ShuffledIterator<T> {
        return ShuffledIterator(elements: _elements)
    }

    private let _elements: [T]

    init(elements: [T]) {
        self._elements = elements
    }

}

let array = Array(0...1000000)

printTimeElapsedWhenRunningCode(title: "Shuffle lazy") {
    _ = Array(LazyShuffledSequence(elements: array).prefix(2))
}

結果はTime elapsed for Shuffle lazy: 3.75410097837448 sと出ました。全てシャッフルする場合と比べて75秒程度速くなっています。

最初の偶数2個を取得するコードではどうでしょうか。

printTimeElapsedWhenRunningCode(title: "Shuffle lazy") {
    _ = Array(LazyShuffledSequence(elements: array).lazy.filter { $0 % 2 == 0 }.prefix(2))
}

Time elapsed for Shuffle lazy: 4.1856529712677 sという結果になりました。これも高速に取得できています。

注)
実行速度計速関数

func printTimeElapsedWhenRunningCode(title:String, operation:()->()) {
    let startTime = CFAbsoluteTimeGetCurrent()
    operation()
    let timeElapsed = CFAbsoluteTimeGetCurrent() - startTime
    print("Time elapsed for \(title): \(timeElapsed) s")
}