How to get a 2D NumPy array with value 1 at indices represented by values in 1D NumPy array (Python)

Question

How to get a 2D np.array with value 1 at indices represented by values in 1D np.array in Python.

Example:

[1, 2, 5, 1, 2]

should be converted to

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

Here you already know the width (shape[2]) value of the new array beforehand.

I can do it manually but is there any way to do it directly using NumPy methods for faster execution? The dimension of my array is quite large and I have to do this for all iteration. Thus, doing this manually for each iteration is quite computationally demanding.

Answer 1

You can create a array with zeros using np.zeros . The shape the array should be (len(1D array), max(1D array)+1) . Then use NumPy's indexing.

idx = [1, 2, 5, 1, 2]
shape = (len(idx), max(idx)+1)
out = np.zeros(shape)
out[np.arange(len(idx)), idx] = 1
print(out)
[[0. 1. 0. 0. 0. 0.]
 [0. 0. 1. 0. 0. 0.]
 [0. 0. 0. 0. 0. 1.]
 [0. 1. 0. 0. 0. 0.]
 [0. 0. 1. 0. 0. 0.]]