numpy数组中随机选择的索引

Question

我有一个二维的numpy数组，里面填满了整数，值是-1、0和+1。我需要随机选择一个不是零的元素，然后计算它周围相邻元素的总和。

有没有办法获取到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 文档]：

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

这个方法会返回一个包含数组的元组，这些数组对应着非零元素的坐标。接下来，你可以画一个索引：

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

然后用这个索引来决定你感兴趣的坐标是哪一个。