如何在NumPy的3D数组中查找指定2D数组的索引

这问题我之前处理类似需求时也琢磨过，其实完全可以沿用你熟悉的np.where思路，只需要调整一下axis参数就能搞定！

核心实现思路 #

原来在2D数组中找1D数组时，你用了np.all(a==b, axis=1)——这里的axis=1是沿着列的方向检查整行是否匹配。放到3D场景里，我们的数组结构是(样本数, 行数, 列数)，要判断每个样本对应的整个2D数组是否和目标b完全一致，只需要把axis设为(1,2)，沿着行和列两个维度做全匹配检查就行。

完整示例代码 #

import numpy as np

# 定义3D数组a
a = np.array([[[1, 0, 0], [0, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 0], [0, 0, 0]]])
# 目标2D数组b
b = np.array([[1, 0, 0], [0, 0, 0], [0, 0, 0]])

# 检查每个2D子数组是否与b完全匹配
matches = np.all(a == b, axis=(1, 2))
# 提取所有匹配的索引
match_indices = np.where(matches)[0]

# 获取第一个匹配的索引（需先确认存在匹配项）
if len(match_indices) > 0:
    first_match_idx = match_indices[0]
    print(first_match_idx)  # 输出：0
else:
    print("未找到匹配的2D数组")

额外说明 #

• 如果3D数组中有多个和b匹配的2D子数组，match_indices会返回所有对应的索引（比如如果a里第0和第2个2D数组都匹配，就会得到array([0,2])）
• 一定要先判断match_indices的长度再取索引，避免没有匹配时出现IndexError

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