Suppose I have the following arrays:
a = np.zeros([4,3])
b = np.asarray([0,1,2,1])
Now how can I set an element on each row in array a to 1 based on the column index indicated by array b? What I need is an array c which looks like below:
[[ 1. 0. 0.]
[ 0. 1. 0.]
[ 0. 0. 1.]
[ 0. 1. 0.]]
Alternatively, is there a way to directly convert array b to array c?
This is a job for advanced indexing:
a[np.arange(a.shape[0]), b] = 1
For arrays A, I, and J where I and J have integer dtype and identical shape, A[I, J] selects all entries A[I[n], J[n]] of A. The assignment then sets these entries to 1.
You can use NumPy broadcasting to directly get c from b by using max of b as the number of columns in the output array, like so -
(b[:,None] == np.arange(b.max()+1)).astype(float)
Sample run -
In [484]: b
Out[484]: array([0, 1, 2, 1])
In [485]: (b[:,None] == np.arange(b.max()+1)).astype(float)
Out[485]:
array([[ 1., 0., 0.],
[ 0., 1., 0.],
[ 0., 0., 1.],
[ 0., 1., 0.]])
