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!

Answer 1

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))