CODEWITHSUNDEEP
pythonnumpy
jiaying chen
jiaying chen

Reputation: 21

How to set 0-1 matrix using a vector of indices using numpy?

Matrix A:

A = np.array([[3, 0, 0, 8, 3],
      [9, 3, 2, 2, 6],
      [5, 5, 4, 2, 8],
      [3, 8, 7, 1, 2], 
      [3, 9, 1, 5, 5]])

Matrix B: values in each row means the index of each row in matrix A.

B = np.array([[1, 2],
      [3, 4],
      [1, 3],
      [0, 1],
      [2, 3]])

We will set values in A whose index are in B to 1, others to 0. Then the result will be:

A = np.array([[0, 1, 1, 0, 0],
      [0, 0, 0, 1, 1],
      [0, 1, 0, 1, 0],
      [1, 1, 0, 0, 0], 
      [0, 0, 1, 1, 0]])

I don't want to use for loop, how can I do it with numpy?

Upvotes: 2

Views: 790

Answers (2)

Matt Eding
Matt Eding

Reputation: 1002

We can index using arrays. For axis0, we just make a range for 0-len(B) to cover each row. Then for axis1, we transpose B to represent all the column indices we want to access.

>>> C = np.zeros_like(A)
>>> C[np.arange(len(B)), B.T] = 1
>>> C
array([[0, 1, 1, 0, 0],
       [0, 0, 0, 1, 1],
       [0, 1, 0, 1, 0],
       [1, 1, 0, 0, 0],
       [0, 0, 1, 1, 0]])

Upvotes: 2

Paul Panzer
Paul Panzer

Reputation: 53029

>>> B = np.array([[1, 2],
...       [3, 4],
...       [1, 3],
...       [0, 1],
...       [2, 3]])

Convenient but a bit wasteful

>>> np.identity(5,int)[B].sum(1)
array([[0, 1, 1, 0, 0],
       [0, 0, 0, 1, 1],
       [0, 1, 0, 1, 0],
       [1, 1, 0, 0, 0],
       [0, 0, 1, 1, 0]])

More economical but also more typing

>>> out = np.zeros((5,5),int)
>>> out[np.c_[:5],B] = 1
>>> out
array([[0, 1, 1, 0, 0],
       [0, 0, 0, 1, 1],
       [0, 1, 0, 1, 0],
       [1, 1, 0, 0, 0],
       [0, 0, 1, 1, 0]])

Upvotes: 0

Related Questions