индекс случайного выбора из массива numpy


У меня есть 2-d numpy массив, заполненный целыми числами [-1, 0, +1]. Мне нужно выбрать случайный элемент, который не равен нулю, и вычислить сумму соседних элементов.

Есть ли способ получить индекс a 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])

Это дает ничью из правильного образца, но мне нужен индекс выбора, чтобы иметь возможность идентифицировать его соседние элементы. Моя другая идея состоит в том, чтобы просто взять образец из индекса, а затем проверить, является ли элемент ненулевым, но это, очевидно, довольно расточительно, когда есть много нулей.

1 2

1 ответ:

Вы можете использовать lattice.nonzero(), чтобы получить расположение ненулевых элементов [ненулевые документы]:

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

, который возвращает кортеж массивов, соответствующих координатам ненулевых элементов. Затем вы рисуете индекс:

>>> np.random.randint(0, len(lnz[0]))
3
И используйте это, чтобы решить, какая координата вас интересует.