math.extras: tweak weighted-randoms add vose alias method

diff --git a/extra/math/extras/extras.factor b/extra/math/extras/extras.factor
index d03da823cc4cea97cad2d63111a6fda2d133fcf3..92676835bca8f58cd02c725ac88aaee3732a9e80 100644 (file)
--- a/extra/math/extras/extras.factor
+++ b/extra/math/extras/extras.factor
@@ -2,12 +2,12 @@
 ! See https://factorcode.org/license.txt for BSD license
 
 USING: accessors arrays assocs byte-arrays combinators
-combinators.short-circuit compression.zlib grouping kernel math
-math.bitwise math.combinatorics math.constants math.functions
-math.order math.primes math.primes.factors math.statistics
-math.vectors namespaces random random.private ranges
-ranges.private sequences sequences.extras sets sorting
-sorting.extras ;
+combinators.short-circuit compression.zlib grouping kernel
+kernel.private math math.bitwise math.combinatorics
+math.constants math.functions math.order math.primes
+math.primes.factors math.statistics math.vectors namespaces
+random random.private ranges ranges.private sequences
+sequences.extras sequences.private sets sorting sorting.extras ;
 
 IN: math.extras
 
@@ -200,10 +200,11 @@ PRIVATE>
     dup sum '[ _ / dup ^ ] map-product ;
 
 : weighted-random ( histogram -- obj )
-    unzip cum-sum [ last >float random ] [ bisect-left ] bi swap nth ;
+    unzip cum-sum [ last >float random ] keep bisect-left swap nth ;
 
-: weighted-randoms ( histogram length -- seq )
-    swap unzip swap [ cum-sum [ last >float random-generator get ] keep ] dip
+: weighted-randoms ( length histogram -- seq )
+    unzip cum-sum swap
+    [ [ last >float random-generator get ] keep ] dip
     '[ _ _ random* _ bisect-left _ nth ] replicate ;
 
 : unique-indices ( seq -- unique indices )
@@ -389,4 +390,36 @@ PRIVATE>
     [ dup 0 < [ neg ] when ] bi@
     [ 1 ] 2dip reduce-evens reduce-odds 2drop ;
 
-
+TUPLE: vose
+    { n fixnum }
+    { items array }
+    { probs array }
+    { alias array } ;
+
+:: <vose> ( dist -- vose )
+    V{ } clone :> small
+    V{ } clone :> large
+    dist assoc-size :> n
+    n f <array> :> alias
+
+    dist unzip dup [ length ] [ sum ] bi / v*n :> ( items probs )
+    probs [ swap 1 < small large ? push ] each-index
+
+    [ small empty? large empty? or ] [
+        small pop :> s
+        large pop :> l
+        l s alias set-nth
+        l dup probs [ s probs nth + 1 - dup ] change-nth
+        1 < small large ? push
+    ] until
+
+    1 large [ probs set-nth ] with each
+    1 small [ probs set-nth ] with each
+
+    n items probs alias vose boa ;
+
+M:: vose random* ( obj rnd -- elt )
+    obj n>> rnd random* { fixnum } declare
+    dup obj probs>> nth-unsafe rnd (random-unit) >=
+    [ obj alias>> nth-unsafe ] unless
+    obj items>> nth-unsafe ;