weighted randomizer

randomizer.cljc

;; port of https://blog.bruce-hill.com/a-faster-weighted-random-choice

(defn weighted-randomizer [weights]
  (let [N                     (count weights)
        avg                   (/ (reduce + weights) N)
        aliases               (volatile! (vec (repeat N [1 nil])))
        {:keys [smalls bigs]} (->> (map-indexed
                                    (fn [idx weight]
                                      (if (< weight avg)
                                        [[idx (/ weight avg)] :smalls]
                                        [[idx (/ weight avg)] :bigs]))
                                    weights)
                                   (group-by peek))]
    (loop [smalls (mapv first smalls)
           bigs   (mapv first bigs)]

      (let [small (first smalls)
            big   (first bigs)]

        (when (and big small)
          (vswap! aliases assoc
                  (nth small 0)
                  [(nth small 1) (nth big 0)])
          (let [new-val [(nth big 0)
                         (- (nth big 1)
                            (- 1 (nth small 1)))]]
            (if (< (nth new-val 1) 1)
              (recur (assoc smalls 0 new-val)
                     (rest bigs))
              (recur (rest smalls)
                     (assoc bigs 0 new-val)))))))

    (fn weighted-random[]
      (let [r            (* N (rand))
            i            (int r)
            [odds alias] (nth @aliases i)]
        (if (> (- r i) odds)
          alias
          i)))))

(def tester (weighted-randomizer [1 2 10]))
(sort-by val (frequencies (repeatedly  5000 (fn [] (tester)))))