Create a 2D numpy array of one-hot vectors from a 1D array

Question

I have a 1D numpy array, let's call it labels . Each element of this array is a digit (0...9). I want to transform labels array into a 2D numpy array of zeros, let's call it y_train . The new 2D array is such that for the element labels[i] , the i ^th row of y_train , has exactly one element at index label[i] equal to 1.0. One way this could be done is by iterating over the labels array. See below

>>> labels = [1, 5, 3, 9, 4]
>>> y_train = np.zeros((5, 10))
>>> for i in range(len(labels)):
        y_train[i][labels[i]] = 1.0
>>> y_train
array([[0., 1., 0., 0., 0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 1., 0., 0., 0., 0.],
       [0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0., 0., 0., 0., 1.],
       [0., 0., 0., 0., 1., 0., 0., 0., 0., 0.]])

This code snippet depicts what I am tryingto achieve. So, my main question is how can I vectorize this operation for faster computation using numpy?

Answer 1

You can simply use this in place of the loop:

y_train[np.arange(len(labels)), labels] = 1

But I recommend using available library methods in sklearn or other packages.