根据搜索条件从Python数组中返回随机元素

Question

抱歉如果这个问题很简单，但我找了一段时间还是没找到一个简单又有效的解决办法。

我有一个二维的Python列表，里面只包含1和0。

比如说：

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

我想随机返回一个值为1的元素的索引。在这个例子中，我希望返回：

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

而且每个索引被选中的概率是一样的。

我可以遍历这个结构中的每个元素，列出所有符合条件的索引，然后用random.choice(list)随机选择一个，但这样做似乎很慢，我觉得应该有更简洁、更符合Python风格的方法来处理这个问题。我可能会处理一个20x20的数组，并且需要多次执行这个操作，所以我希望这个过程尽可能高效。

提前感谢任何帮助和建议！

Answer 1

我会用一个 NumPy 数组来实现这个：

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

如果你一开始就把数据存储在 NumPy 数组里，这种方法可能会比用列表的列表快很多。

如果你想对同一组数据选择多个索引，还有更快的方法。

Answer 2

random.choice 这个功能可以让我们从一个列表中随机选择一个元素。所以我们只需要用一种叫做列表推导的方法，先创建一个包含元素为1的索引的列表，然后再从中随机选择一个。

我们可以使用下面的列表推导：

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

这意味着我们可以这样做：

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

如果你需要多次执行这个操作，可能值得把生成的索引列表保存下来，这样就可以多次从中选择，而不是每次都重新计算这个列表。

Answer 3

我会用列表推导式来生成一个包含元组的列表，这些元组表示数字1的位置，然后用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))