vose.clj

;; i freely admit that i don't holistically grok this; i just transcribed the algorithm into something slightly more clojurish

"
Algorithm: Vose's Alias Method
Initialization:
Create arrays Alias and Prob, each of size n.
Create two worklists, Small and Large.

Multiply each probability by n.

For each scaled probability pi:
If pi<1, add i to Small.
Otherwise (pi≥1), add i to Large.

While Small and Large are not empty: (Large might be emptied first)
Remove the first element from Small; call it l.
Remove the first element from Large; call it g.
Set Prob[l]=pl.
Set Alias[l]=g.
Set pg:=(pg+pl)−1. (This is a more numerically stable option.)
If pg<1, add g to Small.
Otherwise (pg≥1), add g to Large.

While Large is not empty:
Remove the first element from Large; call it g.
Set Prob[g]=1.

While Small is not empty: This is only possible due to numerical instability.
Remove the first element from Small; call it l.
Set Prob[l]=1.

Generation:
Generate a fair die roll from an n-sided die; call the side i.
Flip a biased coin that comes up heads with probability Prob[i].
If the coin comes up \"heads,\" return i.
Otherwise, return Alias[i]."

(defn vose [weights]
  (let [ct (count weights)
        weighted
        (into {}
              (for [[v p] weights]
                [v (* p ct)]))
        [small large]
        (reduce
         (fn [[small large] [i pi]]
           (if (< pi 1)
             [(conj small i) large]
             [small (conj large i)]))
         [[] []]
         weighted)
        [small large prob alias]
        (loop [[l & msmall :as small] small
               [g & mlarge :as large] large
               prob nil
               alias nil
               weighted weighted]
          (if (and l g)
            (let [w2 (update weighted g (partial + (weighted l) -1))]
              (recur
               (cond->> msmall (< (w2 g) 1) (cons g))
               (cond->> mlarge (>= (w2 g) 1) (cons g))
               (assoc prob l (weighted l))
               (assoc alias l g)
               w2))
            [small large prob alias]))
        prob
        (merge prob
               (into {}
                     (map (juxt identity (constantly 1)) (concat large small))))]
    {:n (count weights)
     :opts (vec (map first weights))
     :prob prob
     :alias alias}))

"
Generation:
Generate a fair die roll from an n-sided die; call the side i.
Flip a biased coin that comes up heads with probability Prob[i].
If the coin comes up \"heads,\" return i.
Otherwise, return Alias[i].
"
(defn generate [{:keys [opts n prob alias]}]
  (let [i (rand-int n)
        base (nth opts i)
        pi (get prob base)
        t? (< (rand) pi)]
    (if t? base (get alias base))))