Use np.identity
or np.eye
. You can try something like this with your input i, and the array size s:
np.identity(s)[i:i+1]
For example, print(np.identity(5)[0:1])
will result:
[[ 1. 0. 0. 0. 0. 0. 0. 0. 0. 0.]]
If you are using TensorFlow, you can use tf.one_hot
: https://www.tensorflow.org/api_docs/python/array_ops/slicing_and_joining#one_hot