从 numpy 数组中随机选择的索引

Question

我有一个用整数 [-1, 0, +1] 填充的二维 numpy 数组。我需要从中选择一个不为零的随机元素并计算其相邻元素的总和。

有没有办法获取numpy.random.choice 的索引？

lattice=np.zeros(9,dtype=numpy.int)
lattice[:2]=-1
lattice[2:4]=1
random.shuffle(lattice)
lattice=lattice.reshape((3,3))

random.choice(lattice[lattice!=0])

这会从正确的样本中抽取，但​​我需要选择的索引才能识别其相邻元素。我的另一个想法是只从索引中采样，然后检查元素是否非零，但是当有很多零时，这显然是非常浪费的。

Answer 1

您可以使用lattice.nonzero() 获取非零元素[nonzero docs] 的位置：

>>> lnz = lattice.nonzero()
>>> lnz
(array([0, 0, 1, 1]), array([1, 2, 0, 1]))

返回对应于非零元素坐标的数组元组。然后你画一个索引：

>>> np.random.randint(0, len(lnz[0]))
3

并使用它来决定您感兴趣的坐标。