Given the sparsity of the array that you build, you might want to use Scipy's sparse matrices, which have the advantage of having a small memory footprint:
import numpy
from scipy import sparse
A = numpy.array([0, 0, 1, 2, 1])
k = 3
B = sparse.coo_matrix((numpy.full(len(A), 1, dtype=int), (numpy.arange(len(A)), A)), shape=(len(A), k))
(coo_matrix()
is described in Scipy's documentation). This gives the intended result:
>>> B.todense()
matrix([[ 1., 0., 0.],
[ 1., 0., 0.],
[ 0., 1., 0.],
[ 0., 0., 1.],
[ 0., 1., 0.]])
but with a small memory footprint (if k
is large enough [larger than a few units]). In order to save even more memory, the dtype
above could be made smaller (depending on your exact needs), with dtype=numpy.int8
or even dtype=bool
.