如何利用索引数组为二维NumPy数组的指定元素赋值

这个问题我之前也碰到过，其实核心是要搞懂NumPy的高级索引逻辑～

首先得说说你之前写法的问题：当你用N[:, A] = 1时，:表示选中所有行，而A是一个长度为m的索引数组，这会让NumPy给每一行都选中A里包含的所有列，相当于每一行的这些列都被设为1，最后整个数组自然全是1了，完全偏离了你的需求。

那正确的做法是什么呢？其实很简单，你需要给行和列都提供一一对应的索引数组：

• 行索引：用np.arange(m)生成0到m-1的数组，刚好和A的长度一致，对应每一行的序号
• 列索引：直接用你的A数组

把这两个数组组合起来作为索引，就能精准定位到N[i, A[i]]的位置，然后赋值即可。完整代码如下：

import numpy as np

m = 2000
n = 10
N = np.zeros((m, n))
A = np.random.choice(n, m)

# 正确赋值：精准定位每个N[i, A[i]]
N[np.arange(m), A] = 1

原理解释 #

这里NumPy会执行配对式索引：它会把np.arange(m)和A里的元素一一配对，依次取N[0, A[0]]、N[1, A[1]]……直到N[m-1, A[m-1]]，然后把这些位置统一设为1，完全符合你想要的“每一行只有A[i]对应的列是1”的效果。

你可以打印前几行验证一下：

# 查看前5行的结果
print(N[:5])

会看到每一行都只有一个1，位置正好对应A数组的前5个值。

备注：内容来源于stack exchange，提问作者g00dds