numpy fancy indexing按列提取子矩阵轴顺序异常问题求解

问题原因 #

你观察到的输出差异本质是numpy花式索引的维度广播规则导致的：

• 当你在行维度传入形状为(2,2)的索引数组，列维度留空（默认全选）时，输出形状为索引数组形状 + 列长度 = (2,2) + 3 = (2,2,3)，刚好符合你按行取子矩阵的预期。
• 当你在列维度传入形状为(2,2)的索引数组，行维度用:全选时，输出形状为行长度 + 索引数组形状 = 3 + (2,2) = (3,2,2)，这就是你拿到的非预期结果的形状。

最优实现方案 #

方案1：直接调整维度（最简洁，兼容所有numpy版本）

不需要提前转置原矩阵，只需要把现有结果的前两个维度交换即可：

import numpy as np
A = np.array([[1,2,3],[4,5,6],[7,8,9]])
col_idx = np.array([[0,1],[1,2]])
B = A[:, col_idx].transpose(1,0,2)

执行后B就是你需要的输出：

array([[[1, 2],
        [4, 5],
        [7, 8]],

       [[2, 3],
        [5, 6],
        [8, 9]]])

方案2：滑动窗口视图（性能最优，无额外内存开销，需要numpy≥1.20.0）

如果你的需求是固定取相邻的2列作为子矩阵，直接用滑动窗口工具更直观，也不会产生额外的数组拷贝：

B = np.lib.stride_tricks.sliding_window_view(A, window_shape=2, axis=1)

输出结果和方案1完全一致。

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