Reservoir Sampling in Clojure

I am learning clojure and decided to start out by trying to write a solution to a fairly simple algorithm, reservoir sampling. As I stated, I am learning clojure specifically and problem solving in a functional language in general. Can someone please take a look at my code and critique it on it's "clojureness". Am I using the right idiomatic conventions, is there a way that performs better (and why), formatting, anything really.

(defn sample-seq [size data]
(loop [sample (transient (vec (take size data)))
idx size
data (drop size data)]
(if (empty? data)
(persistent! sample)
(let [rand-num (rand-int idx)
new-sample (if (< rand-num size)
(assoc! sample rand-num (first data))
sample)]
(recur new-sample (inc idx) (rest data))))))
(println (sample-seq 4 [2.0, 4.0, 7.0, 6.0, 3.0, 8.0, 12.0, 9.0, 4.0, 1.0]))

I think your code is pretty readable and looks like idiomatic enough Clojure code¹. So from a readability standpoint your code seems fine. However performing assoc on a vector of length n takes O(log n) time, so your runtime will be in O(n log n) as opposed to O(n), which an imperative implementation would be in.

However there's not much you can do about this other than perhaps using java arrays imperatively, but that would be very unidiomatic Clojure code. And O(n log n) isn't that bad (definitely not as bad as the O(n^2) I incorrectly claimed before).

(Note that my previous note about transients was wrong as assoc! on a transient vector has the same runtime complexity as on a persistent vector).

¹ Usually you'd avoid index-based loops wherever possible, but the nature of the algorithm makes that pretty much impossible here.

• @MGoDave: if this answer was useful, then you should probably up-vote it to thank @sepp2k... – Xavier Nodet Feb 3 '11 at 10:09
• Note that the log factor in clojure's vectors is large: assoc is O(log32 n). For any size of integer that fits in your computer, this is effectively constant. – amalloy May 19 '11 at 23:32

sepp2k's critique is incorrect - assoc on a vector takes O(log n) time not O(n) time because Clojure uses persistent vectors. You should avoid using assoc! because it is Clojure style to avoid needless destructive behavior.

For more on persistent vectors see: http://blog.higher-order.net/2009/02/01/understanding-clojures-persistentvector-implementation/

• Damn, sorry about that. I corrected my answer. – sepp2k Feb 19 '11 at 11:09

It's a bit unfortunate that the reservoir sampling algorithm is funadmentally imperative. If squeezing the last bit of performance out of it is important to you, I'd try using a Java Array internally. (I say try, you'd want to actually time it to see if you'd actually gained any performance.)

The only other performance tip I'd try is to replace idx size with idx (int size) which should guarantee the type internally.

I wouldn't worry about using transients. The Clojure source code uses them all of the time for the exact purpose you have here. Pods are likely to make the code simpler when they drop. Finally, although Kevin L is definitely correct about assoc being an O(log n) operation, I wouldn't guarantee that was also true of transient vectors.