basis/reservoir-sampling/reservoir-sampling.factor

   1 ! Copyright (C) 2020 Doug Coleman.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: accessors kernel math math.functions random sequences ;
   4 IN: reservoir-sampling
   5
   6 : reservoir-sample-iteration ( iteration k obj sampled -- sampled' )
   7     pick over length > [
   8         [ push ] keep 2nip
   9     ] [
  10         roll random roll dupd < [
  11             swap [ set-nth ] keep
  12         ] [
  13            drop nip
  14         ] if
  15     ] if ;
  16
  17 TUPLE: reservoir-sampler iteration k sampled ;
  18 : <reservoir-sampler> ( k -- sampler )
  19     reservoir-sampler new
  20         V{ } clone >>sampled
  21         0 >>iteration
  22         swap >>k ; inline
  23
  24 : reservoir-sample ( obj sampler -- )
  25     [ sampled>> length ] [ k>> ] [ [ 1 + ] change-iteration -rot ] tri < [
  26         sampled>> push
  27     ] [
  28         [ ] [ iteration>> random dup ] [ k>> ] tri < [
  29             swap sampled>> set-nth
  30         ] [
  31             3drop
  32         ] if
  33     ] if ;