# Pythonic way to select list elements with different probability [duplicate]

``````import random
pos = ["A", "B", "C"]
x = random.choice["A", "B", "C"]
``````

This code gives me either "A", "B" or "C" with equal probability. Is there a nice way to express it when you want "A" with 30%, "B" with 40% and "C" with 30% probability?

• Searching found several similar/identical questions here and here – snapshoe Nov 6 '10 at 14:26
• @ma3: One of the weak points of this site is that as old questions tend to be ignored, there's no mechanism or motivation to improve on them. I think my answer is significantly better than at least the higher answers on those questions--havn't read the lower ones--but nobody would ever see it if I posted it on those. – Glenn Maynard Nov 6 '10 at 14:35
• Starting with Python 3.6 you can use `random.choices` to submit a list of weights/probabilities. – Nick Aug 3 '20 at 3:51

Weights define a probability distribution function (pdf). Random numbers from any such pdf can be generated by applying its associated inverse cumulative distribution function to uniform random numbers between 0 and 1.

If Y has a U[0,1] distribution then F⁻¹(Y) is distributed as F. This is used in random number generation using the inverse transform sampling-method.

``````import random
import bisect
import collections

def cdf(weights):
total = sum(weights)
result = []
cumsum = 0
for w in weights:
cumsum += w
result.append(cumsum / total)
return result

def choice(population, weights):
assert len(population) == len(weights)
cdf_vals = cdf(weights)
x = random.random()
idx = bisect.bisect(cdf_vals, x)
return population[idx]

weights=[0.3, 0.4, 0.3]
population = 'ABC'
counts = collections.defaultdict(int)
for i in range(10000):
counts[choice(population, weights)] += 1
print(counts)

# % test.py
# defaultdict(<type 'int'>, {'A': 3066, 'C': 2964, 'B': 3970})
``````

The `choice` function above uses `bisect.bisect`, so selection of a weighted random variable is done in `O(log n)` where `n` is the length of `weights`.

Note that as of version 1.7.0, NumPy has a Cythonized np.random.choice function. For example, this generates 1000 samples from the population `[0,1,2,3]` with weights `[0.1, 0.2, 0.3, 0.4]`:

``````import numpy as np
np.random.choice(4, 1000, p=[0.1, 0.2, 0.3, 0.4])
``````

`np.random.choice` also has a `replace` parameter for sampling with or without replacement.

A theoretically better algorithm is the Alias Method. It builds a table which requires `O(n)` time, but after that, samples can be drawn in `O(1)` time. So, if you need to draw many samples, in theory the Alias Method may be faster. There is a Python implementation of the Walker Alias Method here, and a numpy version here.

Not... so much...

``````pos = ['A'] * 3 + ['B'] * 4 + ['C'] * 3
print random.choice(pos)
``````

or

``````pos = {'A': 3, 'B': 4, 'C': 3}
print random.choice([x for x in pos for y in range(pos[x])])
``````
• If the ratios are from user input this is potentially dangerous, eg. 99.999999%/0.000001%. – Glenn Maynard Nov 6 '10 at 13:46
• Is it possible to somehow incorporate the dictionary into `random.choice()`? – Vladimir Putin Jun 21 '14 at 23:28
• @SomeGuy: You would need to convert the dictionary into a sequence. Like how my second bit of code does. – Ignacio Vazquez-Abrams Jun 21 '14 at 23:54
• This is one of the most elegant answers I've seen, but unfortunately it does not work if the weights are not integers :( – Akshay Sep 13 '14 at 21:49
• Good approach if the weights are given as small integers. If the weights are floats with 3 digit precision, they have to be multiplied by 10^3, creating a big array and inefficient sampling – Protostome Feb 24 '15 at 10:03

Here's a class to expose a bunch of items with relative probabilities, without actually expanding the list:

``````import bisect
class WeightedTuple(object):
"""
>>> p = WeightedTuple({'A': 2, 'B': 1, 'C': 3})
>>> len(p)
6
>>> p, p, p, p, p, p
('A', 'A', 'B', 'C', 'C', 'C')
>>> p[-1], p[-2], p[-3], p[-4], p[-5], p[-6]
('C', 'C', 'C', 'B', 'A', 'A')
>>> p
Traceback (most recent call last):
...
IndexError
>>> p[-7]
Traceback (most recent call last):
...
IndexError
"""
def __init__(self, items):
self.indexes = []
self.items = []
next_index = 0
for key in sorted(items.keys()):
val = items[key]
self.indexes.append(next_index)
self.items.append(key)
next_index += val

self.len = next_index

def __getitem__(self, n):
if n < 0:
n = self.len + n
if n < 0 or n >= self.len:
raise IndexError

idx = bisect.bisect_right(self.indexes, n)
return self.items[idx-1]

def __len__(self):
return self.len
``````

Now, just say:

``````data = WeightedTuple({'A': 30, 'B': 40, 'C': 30})
random.choice(data)
``````
• Note that I only sorted the keys before inserting (rather than just using `items.iteritems`) so it would have deterministic output, which makes the doctests simpler. – Glenn Maynard Nov 6 '10 at 14:05
• Note that I avoided floating-point: if something can be done clearly with integers, it avoids any question of rounding error affecting the output, and it's much easier for the doctest to be comprehensive. Also, this doesn't care what the sum of the values is, eg. it's implicitly normalizing. For example, if you're picking from a list of weighted mirrors and you disable one of them, you don't have to adjust all of the rest so the weights add up to 1. – Glenn Maynard Nov 6 '10 at 14:51
• Note that `sorted(dict.keys())` is a little redundant. `sorted(d)` does the same thing with one list less. – Thomas Wouters Nov 6 '10 at 14:59
• I wrote a similar class based on this one, in Objective-C stackoverflow.com/a/22243637/129202 – Jonny Mar 7 '14 at 7:10

You can also make use this form, which does not create a list arbitrarily big (and can work with either integral or decimal probabilities):

``````pos = [("A", 30), ("B", 40), ("C", 30)]

from random import uniform
def w_choice(seq):
total_prob = sum(item for item in seq)
chosen = random.uniform(0, total_prob)
cumulative = 0
for item, probality in seq:
cumulative += probality
if cumulative > chosen:
return item
``````
• This can be further optimised by using bisect.bisect to identify the option - instead of iterating through all options. This reduces the iteration from O(N) to O(logN) where N is the number of the options. – Nick Aug 3 '20 at 14:01

There are some good solutions offered here, but I would suggest that you look at Eli Bendersky's thorough discussion of this issue, which compares various algorithms to achieve this (with implementations in Python) before choosing one.

Try this:

``````import random
from decimal import Decimal

pos = {'A': Decimal("0.3"), 'B': Decimal("0.4"), 'C': Decimal("0.3")}
choice = random.random()
F_x = 0
for k, p in pos.iteritems():
F_x += p
if choice <= F_x:
x = k
break
``````
• `TypeError: Cannot convert float to Decimal. First convert the float to a string` – Gabi Purcaru Nov 6 '10 at 15:18
{"ft_lang_label":"__label__en","ft_lang_prob":0.7338055,"math_prob":0.97216785,"size":4923,"snap":"2021-04-2021-17","text_gpt3_token_len":1461,"char_repetition_ratio":0.111201465,"word_repetition_ratio":0.11349306,"special_character_ratio":0.33373958,"punctuation_ratio":0.20667285,"nsfw_num_words":0,"has_unicode_error":false,"math_prob_llama3":0.99742544,"pos_list":[0],"im_url_duplicate_count":[null],"WARC_HEADER":"{\"WARC-Type\":\"response\",\"WARC-Date\":\"2021-01-20T13:06:00Z\",\"WARC-Record-ID\":\"<urn:uuid:a4116dbe-dd64-4f2f-91fd-80aa59b56d0d>\",\"Content-Length\":\"200142\",\"Content-Type\":\"application/http; msgtype=response\",\"WARC-Warcinfo-ID\":\"<urn:uuid:423fc3f6-08a6-459b-b01d-c35e99942342>\",\"WARC-Concurrent-To\":\"<urn:uuid:ea89e5b7-ef0d-4df5-8298-7adf32b561dd>\",\"WARC-IP-Address\":\"151.101.193.69\",\"WARC-Target-URI\":\"https://stackoverflow.com/questions/4113307/pythonic-way-to-select-list-elements-with-different-probability\",\"WARC-Payload-Digest\":\"sha1:XCUNVZ7RAPIEIHADO25DKOAMD5BUM6LF\",\"WARC-Block-Digest\":\"sha1:EQXVUE26HJWJFILFDQIPQARPMNTRD2T6\",\"WARC-Identified-Payload-Type\":\"text/html\",\"warc_filename\":\"/cc_download/warc_2021/CC-MAIN-2021-04/CC-MAIN-2021-04_segments_1610703520883.15_warc_CC-MAIN-20210120120242-20210120150242-00573.warc.gz\"}"}