高效计算唯一元素的数量 - NumPy / Python

Question

运行np.unique()时，它会先将数组展平，对其进行排序，然后找到唯一值。当我有形状为（10，3000，3000）的数组时，查找唯一值大约需要1秒钟，但是由于我需要多次调用np.unique()，因此这很快就会累加。由于我只关心数组中唯一数字的总数，排序似乎是浪费时间。

除了np.unique()，是否有更快的方法来查找大数组中唯一值的总数？

- onepint16oz

---

你的数组是什么数据类型（例如 a.dtype 是什么）？ - Warren Weckesser

@WarrenWeckesser uint8 - onepint16oz

Pandas的唯一函数不会排序，因此速度更快。您可能想要查看此链接：https://pandas.pydata.org/pandas-docs/stable/generated/pandas.unique.html - ayhan

@twopint32oz，你能测试一下这两个解决方案的时间吗？看起来它们的时间有些出入。你能同时报告CPU型号和RAM吗？谢谢！ - Divakar

Answer 1

以下是使用 np.uint8 数据类型的数组，比使用 np.unique 更快的方法。

首先，创建要处理的数组：

In [128]: a = np.random.randint(1, 128, size=(10, 3000, 3000)).astype(np.uint8)

为了后续的比较，使用np.unique查找唯一值：

In [129]: u = np.unique(a)

这是更快的方法，v 将包含结果：

In [130]: q = np.zeros(256, dtype=int)

In [131]: q[a.ravel()] = 1

In [132]: v = np.nonzero(q)[0]

验证我们得到了相同的结果：

In [133]: np.array_equal(u, v)
Out[133]: True

定时：

In [134]: %timeit u = np.unique(a)
2.86 s ± 9.02 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [135]: %timeit q = np.zeros(256, dtype=int); q[a.ravel()] = 1; v = np.nonzero(q)
300 ms ± 5.52 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

使用np.unique()耗时2.86秒，而替代方法只需要0.3秒。

- Warren Weckesser

---

这个完美地运行了。谢谢。看起来这也可以用来返回一个排序后的唯一数组。我想知道为什么numpy决定以他们现在的方式实现unique()函数。 - onepint16oz

numpy.unique 可以处理任意数据类型的数组（包括各种整数和浮点数类型、复合类型和结构化数组）。似乎可以通过这种方法将一（或可能两个）字节整数类型作为特殊情况处理，使其能够被 numpy.unique 处理。 - Warren Weckesser

除了计数数组q很快变得巨大之外，没有什么阻止您在更高的位深度上使用这种技术，对吧？它也适用于uint16，只需要0.6秒而不是5.6秒。 - endolith

Answer 2

我们可以利用元素被限制在uint8范围内的事实，通过np.bincount进行分箱计数，然后简单地计算其中非零元素的数量。由于np.bincount需要一个1D数组，我们将输入展平为np.ravel()，然后将其提供给bincount。

因此，实现如下 -

(np.bincount(a.ravel())!=0).sum()

运行时测试

创建包含不同数量唯一数字的输入数组的辅助函数 -

def create_input(n_unique):
    unq_nums = np.random.choice(np.arange(256), n_unique,replace=0)
    return np.random.choice(unq_nums, (10,3000,3000)).astype(np.uint8)

其他方法：

# @Warren Weckesser's soln
def assign_method(a):
    q = np.zeros(256, dtype=int)
    q[a.ravel()] = 1
    return len(np.nonzero(q)[0])

提议方法的验证 -

In [141]: a = create_input(n_unique=120)

In [142]: len(np.unique(a))
Out[142]: 120

In [143]: (np.bincount(a.ravel())!=0).sum()
Out[143]: 120

时间 -

In [124]: a = create_input(n_unique=128)

In [125]: %timeit len(np.unique(a)) # Original soln
     ...: %timeit assign_method(a)  # @Warren Weckesser's soln
     ...: %timeit (np.bincount(a.ravel())!=0).sum()
     ...: 
1 loop, best of 3: 3.09 s per loop
1 loop, best of 3: 394 ms per loop
1 loop, best of 3: 209 ms per loop

In [126]: a = create_input(n_unique=256)

In [127]: %timeit len(np.unique(a)) # Original soln
     ...: %timeit assign_method(a)  # @Warren Weckesser's soln
     ...: %timeit (np.bincount(a.ravel())!=0).sum()
     ...: 
1 loop, best of 3: 3.46 s per loop
1 loop, best of 3: 378 ms per loop
1 loop, best of 3: 212 ms per loop

- Divakar

---

不错。我也尝试过使用 bincount，但它似乎更慢。实际上，当我使用你的代码时，assign_method(a) 花费了 295 毫秒，而 (np.bincount(a.ravel())!=0).sum() 则花费了 425 毫秒。真是难以理解。 - Warren Weckesser

@WarrenWeckesser 可能是硬件问题（CPU、RAM）。我使用的是Intel i7-6700HQ，16GB RAM。但是，你的问题在ideone上可以重现：https://ideone.com/WDWq9j - Divakar

我使用的是“Late 2013”款Macbook Pro，2.6 GHz英特尔Core i7处理器，16 GB 1600 MHz DDR3内存。此外还有：Python 3.5.2，numpy 1.13.1。 - Warren Weckesser

@Divakar 我使用的是2015年款Macbook Pro，2.9 GHz英特尔Core i5处理器，8GB 1867MHz DDR3内存，Python 3.5和numpy 1.12.1。对于128个唯一值，我得到了320毫秒的assign_method(a)和592毫秒的(np.bincount(a.ravel()) !=0).sum()。对于256个唯一值，结果类似。 - onepint16oz

@twopint32oz 感谢您回复了这些！ - Divakar

Answer 3

如果您不介意使用Numba进行JIT编译，并修改您的代码使其易于让Numba发挥其魔力，则可以通过其他答案中已列出的建议，获得一些良好的改进效果。

使用@Divakar帖子中的命名：

from numba import jit
import numpy as np

def create_input(n_unique):
    unq_nums = np.random.choice(np.arange(256), n_unique, replace=0)
    return np.random.choice(unq_nums, (10, 3000, 3000)).astype(np.uint8)

def unique(a):
    return len(np.unique(a))

def assign_method(a):
    q = np.zeros(256, dtype=int)
    q[a.ravel()] = 1
    return len(np.nonzero(q)[0])

def bincount(a):
    return (np.bincount(a.ravel())!=0).sum()

def numba_friendly(a):
    q = np.zeros(256, dtype=int)
    count = 0
    for x in a.ravel():
        if q[x] == 0:
            q[x] = 1
            count += 1
    return count

unique_jit = jit(unique)
assign_method_jit = jit(assign_method)
bincount_jit = jit(bincount)
numba_friendly_jit = jit(numba_friendly)

基准测试：

a = create_input(n_unique=128)

%timeit unique(a)
%timeit unique_jit(a)
%timeit assign_method(a)
%timeit assign_method_jit(a)
%timeit bincount(a)
%timeit bincount_jit(a)
%timeit numba_friendly_jit(a)

结果：

unique:               7.5 s ± 1.14 s per loop (mean ± std. dev. of 7 runs, 1 loop each)
unique_jit:          13.4 s ± 2.03 s per loop (mean ± std. dev. of 7 runs, 1 loop each)
assign_method:       388 ms ± 84.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
assign_method_jit:   341 ms ± 27.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
bincount:            2.71 s ± 218 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
bincount_jit:        138 ms ± 40.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
numba_friendly_jit: 56.4 ms ± 8.96 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

- fuglede

---

你可以修改numba方法来计算唯一元素，对吧？ - endolith

当然可以；只需注意，如果您想将其作为底层结构，则可能需要自己编写哈希映射表；但是，手头的示例足够小，可以将计数放入数组中。 - fuglede