CODEWITHSUNDEEP
python
Difio
Difio

Reputation: 155

Creating a list of random numbers with weights and restrictions

I'd like to create a list of 3 numbers with statistical weights attached. I have found the random.choices function which does that amicably.

random.choices(population=[0,1,2], weights = [0.4,0.3,0.3], k=100)

But in addition to that I'd like to apply some restrictions on this random number generation where I'd like to

random.choices() does not have this functionality it seems. Does anybody have a suggestion on how to approach this?

Kind regards and thanks in advance for your help!

Upvotes: 1

Views: 92

Answers (1)

H. Doebler
H. Doebler

Reputation: 651

Since the samples output by random.choices (remember: with replacement) are statistically independet, just throw out those members that violate your restrictions. Here is an example generator:

def restricted_choices(population, weights=None, k=1, restrictions={}):
    """Yield k values from `population` with restrictions.
    Works almost like `random.choices()`.

    For every item (k, v) in `restrictions` holds:
      the yielded sequence does not include a sub-sequence `v * [k]`.
    """
    N = 0 # count how many values we have yielded so far
    last_value = None # last value that was yielded
    repeat_count = 0 # how often it has been yielded in a row
    while N < k:
        while True:
            x = random.choices(population, weights)[0]
            if x == last_value and 1 + repeat_count == restrictions.get(x, 0):
                continue
            break
        yield x
        N += 1
        if x == last_value:
            repeat_count += 1
        else:
            repeat_count = 1
            last_value = x

Use it like this:

list(restricted_choices(population=[0,1,2], weights = [0.4,0.3,0.3],
                        k=100, restrictions={0:2, 2:3}))

Note however, that the resulting sequence will not be [.4, .3, .3]-distributed any more.

Upvotes: 2

Related Questions