如何高效统计大规模整数标签数组中各标签出现次数并生成同尺寸数组？

问题描述 #

我有一个整数标签数组，希望统计每个标签的出现次数，并将这些次数值存储在与输入尺寸相同的数组中。当前我通过以下循环实现该功能：

import numpy as np

def counter(labels):
    sizes = np.zeros(labels.shape)
    for num in np.unique(labels):
        mask = labels == num
        sizes[mask] = np.count_nonzero(mask)
    return sizes

输入示例：

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

返回结果：

array([[ 2., 5., 1., 4.],
       [ 2., 5., 5., 4.],
       [ 4., 5., 4., 5.]])

但当处理含60000个唯一标签的大规模数组时，该循环耗时过长。这是复杂算法的第一步，我需要将此步骤的耗时控制在30秒以内。请问是否存在可实现该功能的现成函数？若没有，该如何优化现有循环？

解决方案 #

嗨，这个问题我之前也碰到过——Python层面的循环处理几万级别的唯一标签，开销确实大得离谱。好在numpy有几个向量化的方法能完美解决这个问题，效率能提升好几个数量级，完全能满足30秒以内的要求。

方法1：np.unique + np.searchsorted（通用场景） #

这个方法适用于所有整数标签场景，不管标签是否连续：

import numpy as np

def counter_fast(labels):
    # 获取排序后的唯一标签，以及每个标签对应的出现次数
    unique_labels, counts = np.unique(labels, return_counts=True)
    # 用二分查找快速定位每个原标签在unique_labels中的索引，再映射到对应计数
    return counts[np.searchsorted(unique_labels, labels)]

方法2：np.bincount（连续整数标签场景） #

如果你的标签是从0开始的连续整数（没有断层），np.bincount会是更快的选择，它直接统计每个整数的出现次数：

def counter_fast_bincount(labels):
    # 先展平数组统计所有标签的次数
    counts = np.bincount(labels.flatten())
    # 直接用原数组作为索引，获取每个位置对应的次数
    return counts[labels]

为什么这些方法更快？ #

• 原循环的时间复杂度是O(n*k)（n是数组总元素数，k是唯一标签数），而np.unique+searchsorted是O(n log k)，np.bincount是O(n)，当k=6万时，效率提升非常显著。
• 所有核心操作都在numpy的底层C代码中完成，避免了Python循环的额外开销。

验证结果 #

用你提供的输入数组测试，两种方法都会返回和原函数完全一致的结果：

array([[2, 5, 1, 4],
       [2, 5, 5, 4],
       [4, 5, 4, 5]])

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