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

Question

我有一个用整数[-1,0，+ 1]填充的2-d 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

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