Convert a vector to a mask matrix using numpy

Question

Assume we have the following vector:

v = np.array([4, 0, 1])

The goal is to create the 5 x 3 matrix M as follows:

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

Only one element in each column is equal to 1 for the corresponding index in v . For example, since v[0] is 4 then M[4, 0] == 1 , and since v[2] is 1 then M[1, 2] == 1 .

How can I build such a matrix in Python using scipy and numpy? In MATLAB you can do this with the sparse and full functions in a single line. I'd prefer not to use a for loop since I am looking for a vectorized implementation of this.

Answer 1

You can do:

from scipy import sparse

inds = np.array([4, 0, 1])
values = np.ones_like(inds)       # [1, 1, 1]
index = np.arange(inds.shape[0])  # 3
m = sparse.csc_matrix((values, (inds, index)), shape=(5, 3))

Output:

>>> m.todense()
matrix([[0, 1, 0],
        [0, 0, 1],
        [0, 0, 0],
        [0, 0, 0],
        [1, 0, 0]])

Answer 2

If you want a dense array output, you could just use two integer arrays to index the rows/cols of the nonzero elements:

v = np.array([4, 0, 1])
x = np.zeros((5, 3), np.int)
x[v, np.arange(3)] = 1

print(x)
# [[0 1 0]
#  [0 0 1]
#  [0 0 0]
#  [0 0 0]
#  [1 0 0]]