CODEWITHSUNDEEP
pythonarraysnumpyrandom
Artturi Björk
Artturi Björk

Reputation: 3903

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.

Upvotes: 2

Views: 1281

Answers (1)

DSM
DSM

Reputation: 353569

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.

Upvotes: 4

Related Questions