Numpy中基于指定参考数组统计另一数组元素出现频次的方法

问题背景 #

需要统计数组arr1中各元素在已排序、无重复的参考数组arr2中的出现频次，要求包含频次为0的元素的统计值。
示例输入：

arr1 = np.array([1, 0, 3, 0, 3, 0, 3, 0, 8, 0, 1, 8, 0])
arr2 = np.array([0, 1, 2, 8])

预期输出：

freq = np.array([6, 2, 0, 2])

原有实现np.unique(arr1, return_counts=True)只能统计arr1中实际存在的元素，无法输出频次为0的项，以下是可行的实现方案： #

方案1：高效适配排序参考数组（推荐） #

利用arr2已排序的特性，用np.searchsorted匹配位置，再统计频次，时间复杂度低，适合大数据量场景：

import numpy as np

# 匹配arr1元素在arr2中的插入位置
idx = np.searchsorted(arr2, arr1)
# 过滤掉arr1中不存在于arr2的元素：先判断索引不越界，再验证位置对应的元素确实匹配
valid_mask = (idx < len(arr2)) & (arr2[idx] == arr1)
# 统计频次，minlength保证输出长度和arr2一致，自动补全频次为0的项
freq = np.bincount(idx[valid_mask], minlength=len(arr2))

方案2：小数据量直观实现 #

如果数组规模不大，可以用广播对比的方式，代码更简洁易读：

# arr2[:, None]将arr2转为列向量，和arr1广播对比后按行求和得到每个元素的出现次数
freq = (arr1 == arr2[:, None]).sum(axis=1)

方案3：基于原有np.unique的补全方案 #

如果习惯用np.unique，可以对输出结果做补全适配：

unique_vals, counts = np.unique(arr1, return_counts=True)
freq = np.zeros(len(arr2), dtype=int)
# 匹配arr1独有的元素在arr2中的位置
pos = np.searchsorted(arr2, unique_vals)
# 仅给确实存在于arr2中的元素赋值计数
valid = (pos < len(arr2)) & (arr2[pos] == unique_vals)
freq[pos[valid]] = counts[valid]

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