Returning random element from a Python array based on search criteria

Question

Apologies if this is straightforward, but I’ve been looking for a little while now and can’t find a simple, efficient solution.

I have a two-dimensional Python list of lists which only consists of 1’s and 0’s.

e.g.:

a=[[0,1,0],[0,1,1],[1,0,1]]

I wish to return, at random, the indices of a random element which is = 1. In this case I would like to return either:

[0,1], [1,1], [1,2], [2,0], or [2,2]

with an equal probability.

I could iterate through every element in the structure and compile a list of eligible indices and then choose one at random using random.choice(list) – but this seems very slow and I can’t help feeling there is a neater, more Pythonic way to approach this. I will be doing this for probably a 20×20 array and will need to do it many times, so I could do with it being as efficient as possible.

Thanks in advance for any help and advice!

Asked By: Scott

||

Source

Answer 1

When you get your result from random.choice check if it is how you would like it with the correct elements if it is not random again

def return_random(li):
    item = random.choice(li)
    if item == 1: #insert check here
        return item
    else:
        return_random(li)

Edit: to avoid confusion with re module, thanks

Answered By: Jakob Bowyer

Answer 2

I’d use a list comprehension to generate a list of tuples (positions of 1), then random.choice :

from random import choice

a = [[0,1,0],[0,1,1],[1,0,1]]
mylist = []

[[mylist.append((i,j)) for j, x in enumerate(v) if x == 1] for i, v in enumerate(a)]
print(choice(mylist))

Answered By: Vincent Savard

Answer 3

I would use a NumPy array to achieve this:

from numpy import array
random_index = tuple(random.choice(array(array(a).nonzero()).T))

If your store your data in NumPy arrays right from the beginning, this approach will probably be faster than anything you can do with a list of lists.

If you want do choose many indices for the same data, there are even faster approaches.

Answered By: Sven Marnach

Answer 4

random.choice allow us to pick an element at random from a list, so we just need to use a list comprehension to create a list of the indexes where the elements are 1 and then pick one at random.

We can use the follow list comprehension:

>>> a = [[0,1,0],[0,1,1],[1,0,1]]
>>> [(x,y) for x in range(len(a)) for y in range(len(a[x])) if a[x][y] == 1]
[(0, 1), (1, 1), (1, 2), (2, 0), (2, 2)]

Which means we can do:

>>> import random
>>> random.choice([(x,y) for x in range(len(a)) for y in range(len(a[x])) if a[x][y] == 1])
(1, 1)

If you will be doing this many times it may be worth caching the list of indexes generated by the comprehension and then picking from it several times, rather than calculating the list comprehension every single time.

Answered By: David Webb

Answer 5

Another idea would be to store the data in a completely different way: Instead of a list of lists, use a set of index pairs representing the entries that are 1. In your example, this would be

s = set((0, 1), (1, 1), (1, 2), (2, 0), (2, 2))

To randomly choose an index pair, use

random.choice(list(s))

To set an entry to 1, use

s.add((i, j))

To set an entry to 0, use

s.remove((i, j))

To flip an entry, use

s.symmetric_difference_update([(i, j)])

To check if an entry is 1, use

(i, j) in s

Answered By: Sven Marnach

Returning random element from a Python array based on search criteria

Question:

Answers: