CODEWITHSUNDEEP
pythonarraysnumpymatrix-indexing
Alejandro Sazo
Alejandro Sazo

Reputation: 795

How do I use a different index for each row in a numpy array?

I have a NxM numpy array filled with zeros and a 1D numpy array of size N with random integers between 0 to M-1. As you can see the dimension of the array matches the number of rows in the matrix. Each element in the integer array means that at that given position in its corresponding row must be set to 1. For example:

# The matrix to be modified
a = np.zeros((2,10))
# Indices array of size N
indices = np.array([1,4])
# Indexing, the result must be
a = a[at indices per row]
print a

[[0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
 [0, 0, 0, 0, 1, 0, 0, 0, 0, 0]]

I tried using the indexing a[:,indices] but this sets the same indices for each row, and this finally sets all the rows with ones. How can I set the given index to 1 per row?

Upvotes: 7

Views: 8222

Answers (2)

Ersel Er
Ersel Er

Reputation: 771

You should also check the np.eye() function which does exactly what you want. It basically creates 2D arrays filled with zero and diagonal ones. 

>>> np.eye(a.shape[1])[indices]
array([[0., 1., 0., 0., 0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 1., 0., 0., 0., 0., 0.]])

Upvotes: 2

Kasravnd
Kasravnd

Reputation: 107347

Use np.arange(N) in order to address the rows and indices for columns:

>>> a[np.arange(2),indices] = 1
>>> a
array([[ 0.,  1.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  1.,  0.,  0.,  0.,  0.,  0.]])

Or:

>>> a[np.where(indices)+(indices,)] = 1
>>> a
array([[ 0.,  1.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  1.,  0.,  0.,  0.,  0.,  0.]])

Upvotes: 14

Related Questions