Python非对称切片：NumPy矩阵循环索引的广播化实现问询

API 探路者

2026-5-21

用Numpy广播实现目标子数组提取

当然可以用广播来实现这个索引操作，而且相比循环，这种向量化的方式更简洁高效，完全符合Numpy的最佳实践！

先回顾一下我们的基础数据：

import numpy as np
X = np.arange(9).reshape(3,3)
# 原矩阵X：
# array([[0, 1, 2],
#        [3, 4, 5],
#        [6, 7, 8]])

col = [0,1,2]
row = [[0,1],[1,2],[0,1]]

你的循环逻辑是对每一列i，取X[row[i], col[i]]作为结果数组的第i列。要转成广播实现，核心是让row和col的形状能匹配广播规则，从而一次性生成所有需要的索引对。

具体实现代码

# 先把列表转成numpy数组（也可以直接用列表，但数组更灵活）
col_arr = np.array(col)
row_arr = np.array(row)

# 给col增加一个维度，从(3,)变成(3,1)，这样就能和row的(3,2)形状广播匹配
result = X[row_arr, col_arr[:, np.newaxis]].T

print(result)
# 输出正好是目标数组：
# array([[0, 4, 2],
#        [3, 7, 5]])

原理解释

col_arr[:, np.newaxis]（或者用col_arr.reshape(-1,1)）把列索引数组从一维变成二维，形状从(3,)变为(3,1)。
此时row_arr的形状是(3,2)，col_arr[:, np.newaxis]的形状是(3,1)，Numpy会自动广播这两个数组到(3,2)的形状，每个位置的索引对就是(row_arr[i,j], col_arr[i])——这正好对应你循环里第i列的第j个元素的索引。
用这两组索引去取X的元素，得到的是(3,2)的数组，每一行对应原循环里的一列，最后转置.T就得到我们需要的(2,3)目标数组。