Reputation: 5696
I have the following numpy array matrix
,
matrix = np.zeros((3,5), dtype = int)
array([[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0]])
Suppose I have this numpy array indices
as well
indices = np.array([[1,3], [2,4], [0,4]])
array([[1, 3],
[2, 4],
[0, 4]])
Question: How can I assign 1
s to the elements in the matrix
where their indices are specified by the indices
array. A vectorized implementation is expected.
For more clarity, the output should look like:
array([[0, 1, 0, 1, 0], #[1,3] elements are changed
[0, 0, 1, 0, 1], #[2,4] elements are changed
[1, 0, 0, 0, 1]]) #[0,4] elements are changed
Upvotes: 2
Views: 3257
Reputation: 3751
this involves loop and hence may not be very efficient for large arrays
for i in range(len(indices)):
matrix[i,indices[i]] = 1
> matrix
Out[73]:
array([[0, 1, 0, 1, 0],
[0, 0, 1, 0, 1],
[1, 0, 0, 0, 1]])
Upvotes: 0
Reputation: 221524
Here's one approach using NumPy's fancy-indexing
-
matrix[np.arange(matrix.shape[0])[:,None],indices] = 1
Explanation
We create the row indices with np.arange(matrix.shape[0])
-
In [16]: idx = np.arange(matrix.shape[0])
In [17]: idx
Out[17]: array([0, 1, 2])
In [18]: idx.shape
Out[18]: (3,)
The column indices are already given as indices
-
In [19]: indices
Out[19]:
array([[1, 3],
[2, 4],
[0, 4]])
In [20]: indices.shape
Out[20]: (3, 2)
Let's make a schematic diagram of the shapes of row and column indices, idx
and indices
-
idx (row) : 3
indices (col) : 3 x 2
For using the row and column indices for indexing into input array matrix
, we need to make them broadcastable against each other. One way would be to introduce a new axis into idx
, making it 2D
by pushing the elements into the first axis and allowing a singleton dim as the last axis with idx[:,None]
, as shown below -
idx (row) : 3 x 1
indices (col) : 3 x 2
Internally, idx
would be broadcasted, like so -
In [22]: idx[:,None]
Out[22]:
array([[0],
[1],
[2]])
In [23]: indices
Out[23]:
array([[1, 3],
[2, 4],
[0, 4]])
In [24]: np.repeat(idx[:,None],2,axis=1) # indices has length of 2 along cols
Out[24]:
array([[0, 0], # Internally broadcasting would be like this
[1, 1],
[2, 2]])
Thus, the broadcasted elements from idx
would be used as row indices and column indices from indices
for indexing into matrix
for setting elements in it. Since, we had -
idx = np.arange(matrix.shape[0])
,
Thus, we would end up with -
matrix[np.arange(matrix.shape[0])[:,None],indices]
for setting elements.
Upvotes: 8