Add Sequence.randomElement()

The stdlib's randomElement() method extends Collection, since it
operates by selecting a random index. This adds randomElement() as
a Sequence method, specializing the reservoir sampling algorithm
from randomSample(count:) for a count of 1.

Author: natecook1000

Sources/Algorithms/RandomSample.swift

@@ -235,3 +235,61 @@ extension Sequence {
     return randomSample(count: k, using: &g)
   }
 }
+
+//===----------------------------------------------------------------------===//
+// randomElement()
+//===----------------------------------------------------------------------===//
+
+// This method is a single-element specialization of `randomSample(count:)`,
+// and extends the stdlib's `randomElement()` functionality (which is available
+// for collections) down to sequences.
+
+extension Sequence {
+  /// Randomly selects an element from this sequence, using the given generator
+  /// as the source of randomness.
+  ///
+  /// - Parameter rng: The random number generator to use for sampling.
+  /// - Returns: A random element. If the sequence has no elements, this method
+  ///   returns `nil`.
+  ///
+  /// - Complexity: O(*n*), where *n* is the length of the sequence.
+  @inlinable
+  public func randomElement<G: RandomNumberGenerator>(
+    using rng: inout G
+  ) -> Element? {
+    var iterator = makeIterator()
+    guard var result = iterator.next() else {
+      return nil
+    }
+
+    var w = 1.0
+    while true {
+      w *= nextW(k: 1, using: &rng)
+      var offset = nextOffset(w: w, using: &rng)
+      while offset > 0, let _ = iterator.next() {
+        offset -= 1
+      }
+      guard let nextElement = iterator.next()
+        else { break }
+      
+      result = nextElement
+    }
+
+    return result
+  }
+  
+  /// Randomly selects an element from this sequence.
+  ///
+  /// This method is equivalent to calling `randomSample(using:)`, passing in
+  /// the system's default random generator.
+  ///
+  /// - Returns: A random element. If the sequence has no elements, this method
+  ///   returns `nil`.
+  ///
+  /// - Complexity: O(*n*), where *n* is the length of the sequence.
+  @inlinable
+  public func randomElement() -> Element? {
+    var g = SystemRandomNumberGenerator()
+    return randomElement(using: &g)
+  }
+}

Tests/SwiftAlgorithmsTests/RandomSampleTests.swift

@@ -132,4 +132,37 @@ final class RandomSampleTests: XCTestCase {
     almostAllZero = AlmostAllZeroGenerator(seed: 0)
     _ = c.randomSample(count: k, using: &almostAllZero) // must not crash
   }
+  
+  func testSequenceRandomElement() {
+    XCTAssertNil(emptySequence.randomElement())
+    
+    let expectedRange = expectedRange(for: iterations)
+    let randomElements = (0..<(n*iterations)).map { _ in
+      s.randomElement()!
+    }.frequencies
+    XCTAssertEqual(randomElements.count, n)
+    XCTAssert(randomElements.values.allSatisfy { expectedRange.contains($0) })
+    
+    let oneElementSequence = sequence(first: 0, next: { _ in nil })
+    let randomOneElement = (0..<iterations).map { _ in
+      oneElementSequence.randomElement()!
+    }
+    XCTAssert(randomOneElement.allSatisfy { $0 == 0 })
+    
+    let twoElementSequence = sequence(state: [1, 2].makeIterator(), next: { $0.next() })
+    let randomTwoElement = (0..<(2*iterations)).map { _ in
+      twoElementSequence.randomElement()!
+    }.frequencies
+    XCTAssertEqual(randomTwoElement.count, 2)
+    XCTAssert(randomElements.values.allSatisfy { expectedRange.contains($0) })
+  }
+
+  func testSequenceRandomElementRepeatable() {
+    let seed = UInt64.random(in: 0 ... .max)
+    var generator = SplitMix64(seed: seed)
+    let elements1 = (0..<500).map { _ in s.randomElement(using: &generator)! }
+    generator = SplitMix64(seed: seed)
+    let elements2 = (0..<500).map { _ in s.randomElement(using: &generator)! }
+    XCTAssertEqual(elements1, elements2)
+  }
 }