choose.swift

extension Collection {
  /// Chooses at random the given number of elements, using reservoir sampling
  /// aka Algorithm R.
  ///
  /// - Complexity: O(*n*) where *n* is the collection's count.
  func choose<T>(upTo max: Int, using generator: T) -> [Element]
    where T: RandomNumberGenerator
  {
    var i = index(startIndex, offsetBy: max, limitedBy: endIndex) ?? endIndex
    var result = self[..<i].shuffled(using: generator)
    var c = max

    while i < endIndex {
      c += 1
      // Keep item at `i` w/ probability max/c
      let j = Int.random(in: 0..<c, using: generator)
      if j < max {
        result[j] = self[i]
      }
      formIndex(after: &i)
    }

    return result
  }
  
  func choose(upTo max: Int) -> [Element] {
    return choose(upTo: max, using: Random.default)
  }
}

extension RandomAccessCollection where Index: Hashable {
  /// Chooses at random the given number of elements, using Floyd's algorithm.
  ///
  /// - Complexity: O(*m*) where *m* is the lesser of `max` or `count - max`.
  func choose<T>(upTo max: Int, using generator: T) -> [Element]
    where T: RandomNumberGenerator
  {
    var result: Set<Index> = []
    let c = count
    if max >= c {
      return shuffled(using: generator)
    }
    
    let toPick = Swift.min(max, c - max)
    
    for i in (c - toPick) ..< c {
      let r = Int.random(in: 0...i, using: generator)
      let pickedIndex = index(startIndex, offsetBy: r)
      if result.contains(pickedIndex) {
        result.insert(index(startIndex, offsetBy: i))
      } else {
        result.insert(pickedIndex)
      }
    }
    
    // We need to shuffle the result since `Set` order isn't acutally random.
    if toPick == max {
      return result
        .shuffled(using: generator)
        .map({ self[$0] })
    } else {
      return indices
        .filter({ !result.contains($0) })
        .shuffled(using: generator)
        .map({ self[$0] })
    }
  }
  
  func choose(upTo max: Int) -> [Element] {
    return choose(upTo: max, using: Random.default)
  }
}

// Usage:

(1...10).choose(upTo: 0) // []
(1...10).choose(upTo: 1) // [9]
(1...10).choose(upTo: 3) // [5, 3, 8]
(1...10).choose(upTo: 8) // [4, 8, 7, 6, 3, 9, 2, 10]
(1...10).choose(upTo: 1000) // [4, 8, 7, 6, 3, 9, 2, 10]
"abcdefghijklmnopqrstuvwxyz".choose(upTo: 5) // ["a", "b", "z", "j", "e"]
"abcdefghijklmnopqrstuvwxyz".choose(upTo: 35) // ["a", "b", "z", "j", "e"]