如何用Python获取Numpy数组中子数组的所有匹配位置？

给定形状为(s₁,...,sₖ)的Numpy数组C，以及另一个形状为(sⱼ,...,sₖ)（其中j > 1）的数组A，是否存在Python函数可以返回位置列表[p₁,...,pₗ]？其中每个pᵣ的形式为[x₁,...,xⱼ₋₁]，需满足C[x₁,...,xⱼ₋₁] == A。

手动实现代码 #

我尝试了手动嵌套循环的实现方式：

import numpy as np

# C 是形状为 (3, 2, 3) 的数组
C = np.array([
    [[5, 7, 1], [0, 2, 1]],
    [[4, 7, 2], [-1, 2, 2]],
    [[3, 7, 3], [-3, 2, 3]]
])
# A 是形状为 (3) 的数组
A = np.array([-1, 2, 2])
pos_list = []
for i in range(C.shape[0]):
    for j in range(C.shape[1]):
        if np.all(C[i, j] == A):
            pos_list.append([i, j])

print(pos_list)

输出结果 #

[[1, 1]]

更高效的Numpy原生实现 #

不需要手动写循环，利用Numpy的广播和数组操作可以更简洁高效地完成需求：

import numpy as np

C = np.array([
    [[5, 7, 1], [0, 2, 1]],
    [[4, 7, 2], [-1, 2, 2]],
    [[3, 7, 3], [-3, 2, 3]]
])
A = np.array([-1, 2, 2])

# 生成布尔掩码，标记所有匹配A的子数组位置
match_mask = np.all(C == A, axis=tuple(range(len(A.shape), len(C.shape))))
# 获取所有匹配位置的坐标并转为列表
positions = np.argwhere(match_mask).tolist()

print(positions)

代码解释 #

• C == A：Numpy会自动将A广播到与C匹配的形状，逐元素比较得到布尔数组。
• np.all(..., axis=...)：沿着A对应的维度（即从len(A.shape)到C最后一维的轴）取逻辑与，得到形状为(s₁,...,sⱼ₋₁)的布尔掩码，其中True表示对应位置的子数组与A完全匹配。
• np.argwhere(match_mask)：返回掩码中所有True值的坐标，再通过.tolist()转为列表格式，就是需要的位置列表。

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