NumPy数组切片与掩码操作的异常行为问询

问题重现 #

import numpy as np
x = np.empty((2,10,5))
print(x.shape)  # 输出 (2, 10, 5)

print(x[0].shape, x[0,:,:].shape)  # 输出 ((10, 5), (10, 5))

mask = [True,True,True,False,False]
print(x[0,:,mask].shape)  # 输出 (3, 10)

预期结果为(10,3)，但实际得到(3,10)，而二维数组操作符合预期：

y = np.empty((2,5))
print(y.shape)  # 输出 (2, 5)

print(y[0].shape, y[0,:].shape)  # 输出 ((5,), (5,))

print(y[:,mask].shape)  # 输出 (2, 3)

原因解析 #

这不是转置，而是列表型掩码触发了NumPy高级索引的维度重排规则：

• 当混合使用切片（:）和高级索引（列表/数组型索引）时，若高级索引的维度与切片维度不连续，NumPy会将高级索引对应的维度前置。
• 在x[0,:,mask]中：
  • x[0]是(10,5)的二维数组；
  • 第一维用切片:（对应10个元素），第二维用列表掩码（选3个元素）；
  • 这里切片和高级索引作用在不同的非连续维度（对二维数组来说是第一维和第二维），因此结果维度会把高级索引的维度放在前面，得到(3,10)。
• 而二维数组y[:,mask]中，切片作用在第一维，高级索引作用在第二维，两者是连续的维度顺序，因此结果保持(2,3)符合预期。

解决方法 #

方法1：将掩码转为布尔数组

使用NumPy布尔数组作为掩码，而非Python列表，此时会触发基本的布尔索引规则，维度顺序保持不变：

mask = np.array([True,True,True,False,False])
print(x[0,:,mask].shape)  # 输出 (10, 3)

方法2：使用np.take指定轴

通过take方法明确指定索引的轴，避免维度重排：

print(x[0].take([0,1,2], axis=1).shape)  # 输出 (10, 3)

内容的提问来源于stack exchange，提问作者pas-calc