多维NumPy数组整数索引优化及索引行为疑问

1. 更高效的实现方式 #

直接利用广播式高级索引实现完全矢量化操作，无需循环：

import numpy as np

my_set = np.zeros((2, 5), dtype=bool)
init_values = np.array([[1,3,4], [0,1,2]], dtype=np.int64)

# 生成形状为(n_samples, 1)的行索引，与init_values广播匹配
my_set[np.arange(my_set.shape[0])[:, None], init_values] = True

该写法会自动将每个样本i对应的init_values[i]中所有列位置设为True，完全利用NumPy的矢量化计算优势，效率远高于循环方案。

2. np.arange与切片:的行为差异 #

• 切片:属于基本索引，作用是选中整个轴的所有元素，索引时会保留原数组的轴结构，广播时仅扩展维度以匹配其他索引，不会实现逐元素的一一对应。
• np.arange生成的是数组索引（高级索引），与另一个数组索引配合时，会触发广播式逐元素索引：每个行索引会对应到同位置的列索引，精准实现“行i对应列init_values[i,j]”的定位逻辑。

举个直观对比：

• my_set[:, init_values[:,0]]：切片:选中所有2行，init_values[:,0]是[1,0]，因此每行都会取这两个列，得到形状(2,2)的数组。
• my_set[np.arange(2), init_values[:,0]]：np.arange(2)是[0,1]，与init_values[:,0]的[1,0]一一对应，精准选中(0,1)和(1,0)两个位置，得到形状(2,)的结果。

3. 之前方法失败的原因 #

• 方法1：直接用init_values索引：
  NumPy会将二维数组init_values当作行索引数组，即每个元素对应要选取的行号。你的my_set仅2行，但init_values中存在3、4这类大于等于行数的元素，因此触发索引越界错误。
• 方法2：my_set[:, init_values]得到(2,2,3)：
  切片:选中所有2行，init_values是(2,3)的列索引数组，两者广播后会生成(2,2,3)的索引组合——每个行都会对应init_values的所有2行列索引，导致结果形状不符合预期，也无法实现“行i对应自身列索引”的需求。
• 方法3：my_set[:, init_values[:,0]]得到(2,2)：
  init_values[:,0]是(2,)的列索引数组，切片:选中所有行后，每行都会被该数组的所有列索引取值，因此每行返回2个元素，最终得到(2,2)的形状，而非你需要的“每行取对应单个列”的结果。
• 方法4：仅支持单列：
  该写法仅针对init_values的单列做了逐元素索引，若要处理多列必须循环，没有利用NumPy的矢量化优势，数据量增长时效率会急剧下降。

内容的提问来源于Stack Exchange，提问作者Rodrigo Laguna