How can I separate a vector into basis vectors?

I have the basis:

>>> m = 6
>>> basis = np.zeros((m, m))
>>> np.fill_diagonal(basis, 1)
array([[1. 0. 0. 0. 0. 0.]
       [0. 1. 0. 0. 0. 0.]
       [0. 0. 1. 0. 0. 0.]
       [0. 0. 0. 1. 0. 0.]
       [0. 0. 0. 0. 1. 0.]
       [0. 0. 0. 0. 0. 1.]]

How can I access the superposition of a vector? Let's say from:

vec = [0, 1, 1, 0, 1, 0]

So I want the output:

array([[0. 1. 0. 0. 0. 0.]
       [0. 0. 1. 0. 0. 0.]
       [0. 0. 0. 0. 1. 0.]]

Solution

We can construct a diagonal with the number of elements of vec, and then filter the rows, so:

>>> vec = np.array([0,1,1,0,1,0])
>>> np.identity(6)[vec.astype(bool)]
array([[0., 1., 0., 0., 0., 0.],
       [0., 0., 1., 0., 0., 0.],
       [0., 0., 0., 0., 1., 0.]])

This is however not very efficient if the number of items in the vec is large, since constructing an identity matrix takes quadratic time in the length of vec.

We can also work with .diag(…) as @AlexAlex says:

>>> vec = np.array([0,1,1,0,1,0])
>>> np.diag(vec)[vec.astype(bool)]
array([[0, 1, 0, 0, 0, 0],
       [0, 0, 1, 0, 0, 0],
       [0, 0, 0, 0, 1, 0]])

Here we can specify different values, for example:

>>> vec = np.array([0,1,2,0,3,0])
>>> np.diag(vec)[vec.astype(bool)]
array([[0, 1, 0, 0, 0, 0],
       [0, 0, 2, 0, 0, 0],
       [0, 0, 0, 0, 3, 0]])