Numpy多维索引：如何按行依据索引数组提取元素？

假设你的需求是对fp的每一行，按照ix中的列索引提取对应元素（最终得到一个(4000, 3, 3)的数组），或者是用ix的每一行对应fp的对应行提取元素（得到(3,3)的数组），下面分别给出两种场景的解决方案：

场景1：对fp的所有行应用ix的索引 #

如果需要给fp的4000行都按照ix的3x3列索引提取元素，直接用NumPy的广播索引即可，代码简洁高效：

import numpy as np
np.random.seed(100)
fp = np.random.rand(4000, 5)
ix = np.random.randint(0, 5, (3, 3))

# 提取结果，形状为(4000, 3, 3)
result = fp[:, ix]

验证示例 #

用你给出的fp前3行测试：

fp_sample = np.array([
    [0.54340494, 0.27836939, 0.42451759, 0.84477613, 0.00471886],
    [0.12156912, 0.67074908, 0.82585276, 0.13670659, 0.57509333],
    [0.89132195, 0.20920212, 0.18532822, 0.10837689, 0.21969749]
])
ix = np.array([[3,4,4],[1,3,4],[4,3,3]])

# 对这3行应用ix索引
sample_result = fp_sample[:, ix]
print(sample_result)

输出：

array([[0.84477613, 0.00471886, 0.00471886],
       [0.67074908, 0.13670659, 0.57509333],
       [0.21969749, 0.10837689, 0.10837689]])

场景2：仅用ix的行对应fp的对应行提取 #

如果你的需求是ix的第i行对应fp的第i行（比如ix是3行，对应fp的前3行），这里有两种可读性强的实现方式：

方法1：使用np.take_along_axis #

这个函数专门用于按指定轴提取索引对应的元素，语义清晰：

# 先将ix转换为和fp前3行匹配的形状（添加行维度）
ix_expanded = ix[:, np.newaxis]
# 提取fp的前3行对应ix的元素，结果形状(3,3)
result = np.take_along_axis(fp[:3], ix_expanded, axis=1).squeeze()

方法2：直接使用高级索引 #

利用NumPy的高级索引特性，直接匹配行和列的对应关系：

# 生成行索引数组，对应ix的每一行
row_indices = np.arange(ix.shape[0])
# 提取对应元素，结果形状(3,3)
result = fp[row_indices, ix]

用示例测试的话，两种方法都会得到和场景1中sample_result完全一致的输出。

关键原理说明 #

• 广播索引：当你用fp[:, ix]时，NumPy会自动将形状为(3,3)的ix广播为(4000,3,3)，从而对fp的每一行都应用相同的列索引规则。
• 高级索引：fp[row_indices, ix]中，row_indices（形状(3,))和ix（形状(3,3))会被广播为相同的(3,3)形状，最终提取fp[row, col]位置的所有元素。

内容的提问来源于stack exchange，提问作者Eric B