如何用无循环方法统计1D数组与3D数组各最后维度的公共元素数

解决方案代码 #

import numpy as np

arr1 = np.array([1, 2, 3])
arr2 = np.array([
    [[1, 6], [2, 7], [3, 8]],
    [[4, 9], [5, 10], [6, 11]],
    [[7, 12], [8, 13], [9, 14]]
])

# 生成布尔掩码：标记arr2中每个元素是否在arr1内
element_in_arr1 = np.isin(arr2, arr1)
# 沿最后维度求和，得到每个子数组的公共元素数量
common_counts = element_in_arr1.sum(axis=-1)

print(common_counts)

代码说明 #

1. 布尔掩码生成：np.isin(arr2, arr1) 会遍历arr2的所有元素，判断其是否存在于arr1中，输出一个与arr2形状完全一致的布尔数组——匹配的元素为True（等价于数值1），不匹配为False（等价于数值0）。
2. 维度求和：sum(axis=-1) 指定沿数组的最后一个维度（即需要单独对比的每个子数组，比如示例中的[2,7]）进行求和，直接得到每个子数组与arr1的公共元素数量。

运行结果 #

[[1 1 1]
 [0 0 0]
 [0 0 0]]

结果对应arr2中每个最后维度的统计值：比如子数组[2,7]与arr1的公共元素数量为1，完全符合需求。

优势 #

全程使用numpy的向量化操作，无需手动编写循环，既保证代码简洁性，又能利用numpy底层的优化实现，处理大规模数组时效率远高于Python循环。

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