为何用[[a,b],[c,d],:]索引(A,B,C)型numpy数组得到(2,C)而非(2,2,C)？

当对形状为(A, B, C)的NumPy数组使用[[a, b], [c, d], :]进行索引时（满足0 <= a, b < A，0 <= c, d < B），得到的结果形状为(2, C)，但分步索引x[[a, b]][:, [c, d]]却能得到预期的(2, 2, C)，核心原因是NumPy花式索引的配对广播规则：

关键差异点 #

• 单次多维花式索引：x[[a, b], [c, d], :]中，第一个维度的索引[a, b]和第二个维度的[c, d]都是长度为2的一维数组。NumPy会将它们按位置一一配对，提取(a, c, :)和(b, d, :)这两个位置的子数组，最终拼接成(2, C)的形状。
• 分步索引：x[[a, b]]先提取第一个维度的a、b行，得到形状为(2, B, C)的数组；再对这个数组的第二个维度用[:, [c, d]]索引，此时是对(2, B, C)的第二个维度做花式索引，会为每一行（共2行）提取c、d列，最终得到(2, 2, C)的笛卡尔积结果。

如何单次索引得到(2, 2, C) #

如果想在单次索引中实现笛卡尔积效果，需要通过调整索引数组的形状触发广播，让两个索引数组扩展为(2, 2)的形状：

import numpy as np

A, B, C = 7, 40, 57
x = np.random.random((A, B, C))
a, b = 4, 1
c, d = 10, 5

# 方法1：将第一个索引转为二维
print(x[[[a], [b]], [c, d], :].shape)  # 输出 (2, 2, 57)
# 方法2：将第二个索引转为二维
print(x[[a, b], [[c], [d]], :].shape)  # 输出 (2, 2, 57)

原示例代码及运行结果 #

代码 #

import numpy as np

A = np.random.randint(100)
B = np.random.randint(100)
C = np.random.randint(100)
print(f"{A=}, {B=}, {C=}")

x = np.random.random((A, B, C))
print(f"{x.shape=}")

a = np.random.randint(0, A)
b = np.random.randint(0, A)
print(f"{a=}, {b=}")

c = np.random.randint(0, B)
d = np.random.randint(0, B)
print(f"{c=}, {d=}")

print(f"{x[[a, b], [c, d], :].shape=}")
print(f"{x[[a, b]][:, [c, d]].shape=}")

运行结果 #

A=7, B=40, C=57
x.shape=(7, 40, 57)
a=4, b=1
c=10, d=5
x[[a, b], [c, d], :].shape=(2, 57)
x[[a, b]][:, [c, d]].shape=(2, 2, 57)

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