CODEWITHSUNDEEP
pythonnumpy
Joseph D.
Joseph D.

Reputation: 12174

Generate a 2D list based on probability

I want to generate an n-row list where each row is k columns.
Then order is based on probability.

Given: L = [1, 2, 3] and the probability list P = [0.3, 0.2, 0.5],

Expected output based on the probability:

[[3, 3],[3, 1],[1, 3],[3, 2],[2, 3],[1, 1]]

where [3,3] is the first since it yields the highest probability of 1.0.
second is [3,1] with 0.8 and so on.

Here's my code:

L = [3, 1, 2]        # Input List
P = [0.3, 0.2, 0.5]  # Probability List
k = 2                # 2 columns per row
n = 6                # 6 rows in 2D array

nL = np.array(L)

generated_list = np.zeros(shape=(1, k), dtype=int)
while generated_list.shape[0] != n:
    choice = np.random.choice(L, size=k, p=P)
    if not np.any(generated_list == [choice]):
        generated_list = np.append(generated_list, [choice], axis=0)

I can't seem to have an ordered list based on probability.
It also contains duplicate elements.

Sample generated_list output:

[[0, 0],[3, 1],[2, 3],[1, 2],[2, 3],[2, 3]]

What's the best way to do this:
1. Compare if a 1-D element is already in 2-D array
2. Have the random choice return first the highest probability combination

Upvotes: 2

Views: 361

Answers (2)

zipa
zipa

Reputation: 27869

This should do it, with a little help of itertools.product:

from itertools import product

L = [1, 2, 3]
P = [0.3, 0.2, 0.5]

n = 2
k = 6

probs = dict(zip(L, P))

result = sorted(list(product(L, repeat=n)), key=lambda x: (sum(probs[i] for i in x), [probs[i] for i in x]), reverse=True)[:k]

Upvotes: 2

ted
ted

Reputation: 14724

You can use sorted:

l = [[0, 0],[3, 1],[2, 3],[1, 2],[2, 3],[2, 3]]
p = [0.1, 0.2, 0.2, 0.5]
sorted(l, key= lambda x: sum(p[v] for v in x))

returns:

[[3, 1], [2, 3], [2, 3], [2, 3], [1, 2], [0, 0]]

(I assumed that you had 4 different probabilities as your generated_list has indexes 0 to 3)

Upvotes: 1

Related Questions