python - 如何为 numpy 数组中的每一行使用不同的索引？

Question

我有一个用零填充的 NxM numpy 数组和一个大小为 N 的一维 numpy 数组，其中随机整数介于 0 到 M-1 之间。如您所见，数组的维度与矩阵中的行数相匹配。整数数组中的每个元素表示在其对应行中的给定位置必须设置为 1。例如:

# The matrix to be modified
a = np.zeros((2,10))
# Indices array of size N
indices = np.array([1,4])
# Indexing, the result must be
a = a[at indices per row]
print a

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

我尝试使用索引 a[:,indices] 但这为每一行设置了相同的索引，这最终将所有行设置为 1。如何将给定索引设置为 每行 1？

Answer 1

使用 np.arange(N) 来处理行和列的索引:

>>> a[np.arange(2),indices] = 1
>>> a
array([[ 0.,  1.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  1.,  0.,  0.,  0.,  0.,  0.]])

或者:

>>> a[np.where(indices)+(indices,)] = 1
>>> a
array([[ 0.,  1.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  1.,  0.,  0.,  0.,  0.,  0.]])