如何用NumPy中两个一维数组的所有组合索引二维数组？

要实现两个一维索引数组的所有组合索引，核心是利用NumPy的广播机制，将其中一个数组调整形状后与另一个数组匹配，从而自动生成所有索引对。

具体实现步骤 #

1. 定义原始数组和索引数组：

import numpy as np

weights = np.array(
    [
        [1, 2, 3],
        [4, 5, 6],
        [7, 8, 9],
    ]
)
positions1 = np.array([2, 1, 0])
positions2 = np.array([1, 1, 0])

2. 将positions1转换为列向量（形状变为(3,1)），再与positions2一起作为索引传入：

# 方法1：用np.newaxis扩展维度
result = weights[positions1[:, np.newaxis], positions2]

# 方法2：用reshape更直观
result = weights[positions1.reshape(-1, 1), positions2]

原理说明 #

• positions1[:, np.newaxis]会把原本形状为(3,)的一维数组转为(3,1)的列向量，此时它和形状为(3,)的positions2满足广播条件。
• NumPy会自动将两者扩展为(3,3)的二维数组：
  • 扩展后的positions1矩阵：[[2,2,2], [1,1,1], [0,0,0]]
  • 扩展后的positions2矩阵：[[1,1,0], [1,1,0], [1,1,0]]
• 两个矩阵对应位置的元素组成所有需要的索引对，最终从weights中提取出目标3×3矩阵。

结果验证 #

运行上述代码后，result的值为：

[[8 8 7]
 [5 5 4]
 [2 2 1]]

完全符合需求的索引组合结果。

内容的提问来源于stack exchange，提问作者J.N.