I have a Cython function that takes a list of weights/probabilities (double) and returns a random index into the list.

For example:

choose_one(np.array([0.1, 0.4, 0.2, 0.3]), 4)

Returns 0 with probability 0.1, 1 with probability 0.4 etc.

The Cython code is:

cdef extern from "stdlib.h":
    double drand48()
    void srand48(long int seedval)

def choose_one(np.ndarray[double, ndim=1] weights, Py_ssize_t length):
    cdef Py_ssize_t idx, i
    cdef double cs
    cdef double random
    random = drand48()
    cs = 0.0
    i = 0
    while cs < random and i < length:
        cs += weights[i]
        i += 1
    return i - 1

I profiled some complex simulation code and found that about 50% of the execution time is spent in this function which is called approximately 3.5 million times. Because it is such a time hog, and such an important utility, I'd like to speed it up if possible.

One of the problems is that the input arrays are often very large with thousands of entries.

I tried various methods from Numpy in Cython and none of them come close to as fast as what I have written in Cython. I am a Cython newbie and am curious if there is something else I need to do to speed this up.

  • \$\begingroup\$ An important bit of information missing is whether the function is called repeatedly with the same weights, or often with the same weights. In this case, a two step implementation, where you first prepare the distribution, and then sample could yield large benefits. \$\endgroup\$
    – Dan Getz
    Jun 18, 2017 at 11:36

2 Answers 2


The performance can be improved a bit more by declaring the function using cpdef (or cdef) instead of def (when the calling code is using Cython types too, check out http://docs.cython.org/src/userguide/language_basics.html#python-functions-vs-c-functions for more info).

Another idea that might help is to check whether the random number is greater than 0.5 and, if it is, count backwards from the end of the array rather than the front (assuming that the values in the array always add up to one).

Also, if this function takes up half of the program's running time then it seems likely that the program could be restructured to use a different algorithm which calls this function less often, but this would probably require much more effort than the other suggestions.

  • \$\begingroup\$ cpdef isn't necessarily faster though it seems, in a quick test it basically doubled the runtime for me. \$\endgroup\$
    – ferada
    Jun 23, 2016 at 13:20
  • 1
    \$\begingroup\$ @ferada It will be faster when called from other Cython code, not necessarily from python. The main reason for the change is to avoid unnecessary conversions between C types and Python objects, but this will only happen if the calling code if using C types too. Actually, I should probably add that caveat to the answer... \$\endgroup\$
    – gfv
    Jun 23, 2016 at 13:42
  • \$\begingroup\$ @gfv Anecdotal evidence: I changed def to cpdef and got about a 32% improvement. \$\endgroup\$ Jun 24, 2016 at 3:46
  • \$\begingroup\$ @RyanRosario Good to hear! \$\endgroup\$
    – gfv
    Jun 27, 2016 at 9:48

If you understand enough C I'd suggest to look at the generated C code to judge whether it's not "optimal" yet. That said it looks like you've done everything for this single function; apart from switching to float instead of double for comparison I don't see anything in particular to change.

With that many invocations though the bottleneck will probably be the switch between Python and C, considering that the function itself does very little you'll not get very far by speeding up the function, instead you'll have to either do more on the C side, i.e. generate multiple results and pass them back (in which you should compare with numpy.random.choice I guess) or pass in more data; besides that you still have the option to parallelise this operation.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.