5 Levels of Generating Weighted Random Numbers in Python | by Yang Zhou | Towards Data Science

To get a random number in Python is simple, since there is a built-in module called random. But how about weighted random numbers?
Let’s see a real interview question:
There is a dictionary as following:
my_nums = {'A': 5, 'B': 2, 'C': 2, 'D': 1}
The keys are four letters and the values are their weights.
Please design a function that can generate a letter randomly based on the weights.
(In other words, this function should have 50% probability to generate “A”, 20% probability to generate “B”, 20% probability to generate “C”, and 10% probability to generate “D”.)
If you are very familiar with Python, you may know that we can solve the problem with a method called random.choices() introduced from Python 3.6, or a similar method in Numpy.
However, a technical interview is not to test if you can remember a special API or not. It aims to test whether you can design and implement an algorithm by yourself. So, the interviewer probably will ask you:
“Can you design and implement the method by yourself?”
In fact, this is a very good interview question. Because programmers of different levels will have different ideas for it. So it will be easy to distinguish between junior and senior Python developers.
This article will dive into 5 different solutions for this problem, from the most intuitive but inefficient way to the best answer.
If our candidate list contains the expected number of every letter, it will become a simple problem.
As shown above, the most intuitive solution is to generate a list which includes the expected number of every letter based on their weights, then just choose randomly from this list.
However, if there are many letters and their weights are bigger numbers, this solution is obviously not very efficient. Cause we will waste lots of time to generate the list.
Actually, the previous solution’s list is not necessary.
Since the total weight is 10 (5+2+2+1). We can just generate a random integer between 1 and 10 firstly, then return a letter based on this number:
The above code avoids generating a list like the previous solution, so it’s more efficient.
But we have to write much if-elif-else code, this looks ugly.
In fact, we don’t have to compare the r with all ranges.
As the following example, we can just subtract one weight from r in turn, anytime the result is less than or equal to zero, we return it:
Obviously, as to the previous solution, the faster r reaches 0, the more efficient our algorithm will be.
So, how can we make the r reach 0 faster?
Intuitively, if the r always subtract the current largest weight, it will reach 0 more quickly. So, before running the weighted_random() function, we can sort my_nums according to the weights from large to small.
If you know the mathematical expectation, we can prove this idea by it.
For example, the follows are my_nums dictionaries with different item orders. For each dictionary, let’s calculate its mathematical expectation about how many steps needed to get a random letter.
{A:5, B:2, C:2, D:1}
E = 5/10*1 + 2/10*2 + 2/10*3 + 1/10*4 = 19/10
{B:2, C:2, A:5, D:1}
E = 2/10*1+ 2/10*2 + 5 /10*3 + 1/10*4 = 25/10
{ D:1, B:2, C:2, A:5}
E = 1/10*1 + 2/10*2 + 2/10*2 + 5/10*5 = 34/10
As shown above, the best order, which sorts weights from large to small, needs the least average steps to get the result.
Our solution is good enough so far. But there still are spaces to make it better.
The level 4 solution, in fact, introduces a new time-consuming step. Cause we have to sort the my_nums dictionary at first. When this is a larger dictionary, it’s totally not efficient enough.
Under the hood, the built-in Python function random.choices() applies a smarter idea. It uses the cumulative weights instead of the original weights. And since the cumulative weights are in ascending order, we can use binary search to find the location faster.
Based on our problem, let’s implement this great idea:
Now, this solution applies the same idea of Python’s random.choices() method. (By the way, we can check the source code of this built-in method on GitHub.)
If we can design our solution at this level, we definitely will ace the technical interview.
To ace technical interviews, just remembering a special built-in method is not enough. You need to know how it works under the hood and it will be great if you can implement it by yourself.
Thanks for reading. If you like it, please follow me and become a Medium member to enjoy more great articles about programming and technologies!
More interesting Python posts for you: