Picking a random state from a set based on probabilities I'm writing a Markov Chain implementation, and need to be able to pick a random state from a set of states, where each state has a different probability of being chosen. In my caffeine-high, I wrote this:

(defn random-state
  "Accepts a map mapping a state to the frequency that it will be chosen.
  Assumes that the frequencies add up to 100%. Will have undesired behavior if that isn't the case."
  [state-frequency-map]
  (let [r (rand)]
    (loop [[[state prob] &amp; rest-states] (seq state-frequency-map)
           prob-sum 0
           last-sum 0]

      (let [new-sum (+ prob-sum prob)]
        (if (&lt;= last-sum r new-sum)
          state
          (recur rest-states new-sum prob-sum))))))


I had very little idea of how I was going to achieve this going in, so I kind of cobbled it together. Basically, it works by maintaining a sum of probabilities, and checking if a randomly generated number (between 0 and 1) falls between the last and current probability sums. For example:


The probability map is {:a 0.3, :b 0.3, :c 0.4}.
The randomly drawn number is 0.7.
0.7 is greater than the sums of :a and :b's probabilities (0.6), but less than :c's, so :c is chosen this time.


This works, but like I said, I kind of just threw it together as I figured out what I was going to do, and now I'm in mind-lock and can't think of improvements for it. Is this optimal? Are there any ways of making this more idiomatic via better use of built-ins? I'm looking for any suggestions here. I'd also be open to suggestions of entirely different approaches of solving this.
A solution that isn't O(n) with respect to the number of states would be nice.

Example of its use showing that it does indeed pick the states with the correct frequencies:

(let [tries 1e7]
  (mapv (fn [[s f]] [s (/ f tries)]) ; Get the percentage occurrence 
    (frequencies
      (for [n (range tries)]
        (random-state
          {\a 0.2
           \b 0.2
           \c 0.5
           \d 0.1})))))

=&gt; [[\b 0.2000032] [\a 0.2000432] [\c 0.4999343] [\d 0.1000193]]
There could potentially be millions of states, since each distinct word in the text I scan will be become a state. And they'll never change once calculated. I'm going to scan over a book, calculate the probabilities, then serialize the state/probability map to use later.


Given this, there are a couple more efficient approaches, depending on how much preprocessing you want to do. There's a great article, Darts, Dice, and Coins, that outlines this further and ends with an O(1) algorithm, Vose's alias method.

I think the most straightforward, though, is a binary search tree-based approach ("roulette wheel selection"). Same basic idea as your current code, fill [0, 1) with proportional sections corresponding to your state probability map:

[  :a  | :b |   :c   |:d|  :e  )
0    0.25  0.4     0.76 0.8    1


but then instead of looping over a list, use something like a TreeMap that is designed to answer these kinds of queries in O(log n).

(defn states-&gt;treemap
  "Builds a navigable map of `partial sum -&gt; state`"
  [states]
  (-&gt;&gt; states
    (reduce (fn [[entries sum] [k v]]
              (let [sum' (+ sum v)]
                [(conj entries [sum' k]) sum']))
            [[] 0])
    first
    (into {})
    (java.util.TreeMap.)))

(defn random-state [^java.util.NavigableMap nm]
  (.getValue
    (.ceilingEntry nm (rand))))


That's not "pure Clojure", obviously, but we're on the JVM and it's pretty idiomatic to pull from Java as needed. Otherwise, your Clojure looks fine to me, but definitely check out that article for further depth.