index of a random choice from numpy array

问题

I have a 2-d numpy array populated with integers [-1, 0, +1]. I need to choose a random element that is not zero from it and calculate the sum of its adjacent elements.

Is there a way to get the index of 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])

This gives the draw from the right sample, but I would need the index of the choice to be able to identify its adjacent elements. My other idea is to just sample from the index and then check if the element is non-zero, but this is obviously quite wasteful when there are a lot of zeros.

回答1:

You can use lattice.nonzero() to get the locations of the nonzero elements [nonzero docs]:

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

which returns a tuple of arrays corresponding to the coordinates of the nonzero elements. Then you draw an index:

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

and use that to decide which coordinate you're interested in.

来源：https://stackoverflow.com/questions/22440102/index-of-a-random-choice-from-numpy-array