使用Cython在数组中查找所有唯一元素的最快方法

Question

我正试图找到最有效的方法来从NumPy数组中查找唯一的值。NumPy的unique函数非常慢，并在查找唯一值之前对值进行排序。Pandas使用klib C库对值进行哈希处理，速度更快。我正在寻找一个Cython解决方案。

最简单的解决方案似乎就是遍历数组并使用Python集合添加每个元素，如下所示：

from numpy cimport ndarray
from cpython cimport set

@cython.wraparound(False)
@cython.boundscheck(False)
def unique_cython_int(ndarray[np.int64_t] a):
    cdef int i
    cdef int n = len(a)
    cdef set s = set()
    for i in range(n):
        s.add(a[i])
    return s

我也尝试了C++的unordered_set

from libcpp.unordered_set cimport unordered_set
@cython.wraparound(False)
@cython.boundscheck(False)
def unique_cpp_int(ndarray[np.int64_t] a):
    cdef int i
    cdef int n = len(a)
    cdef unordered_set[int] s
    for i in range(n):
        s.insert(a[i])
    return s

性能表现

# create array of 1,000,000
a = np.random.randint(0, 50, 1000000)

# Pure Python
%timeit set(a)
86.4 ms ± 2.58 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

# Convert to list first
a_list = a.tolist()
%timeit set(a_list)
10.2 ms ± 74.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

# NumPy
%timeit np.unique(a)
32 ms ± 1.17 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

# Pandas
%timeit pd.unique(a)
5.3 ms ± 257 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

# Cython
%timeit unique_cython_int(a)
13.4 ms ± 1.02 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)

# Cython - c++ unordered_set
%timeit unique_cpp_int(a)
17.8 ms ± 158 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

讨论

因此，pandas比经过Cython优化的集合快约2.5倍。当有更多不同元素时，它的领先优势会增加。令人惊讶的是，纯Python集合（在列表上）击败了Cython集合。

我的问题是，有没有比仅重复使用add方法更快的Cython方法？ C++的unordered_set是否可以改进？

使用Unicode字符串

当我们使用Unicode字符串时，情况就发生了变化。我认为我必须将numpy数组转换为object数据类型，以便为Cython正确添加其类型。

@cython.wraparound(False)
@cython.boundscheck(False)
def unique_cython_str(ndarray[object] a):
    cdef int i
    cdef int n = len(a)
    cdef set s = set()
    for i in range(n):
        s.add(a[i])
    return s

我又尝试了来自c ++的unordered_set

@cython.wraparound(False)
@cython.boundscheck(False)
def unique_cpp_str(ndarray[object] a):
    cdef int i
    cdef int n = len(a)
    cdef unordered_set[string] s
    for i in range(n):
        s.insert(a[i])
    return s

性能

创建一个包含 1 百万个字符串的数组，其中有 1,000 个不同的值。

s_1000 = []
for i in range(1000):
    s = np.random.choice(list('abcdef'), np.random.randint(5, 50))
    s_1000.append(''.join(s))

s_all = np.random.choice(s_1000, 1000000)

# s_all has numpy unicode as its data type. Must convert to object
s_unicode_obj = s_all.astype('O')

# c++ does not easily handle unicode. Convert to bytes and then to object
s_bytes_obj = s_all.astype('S').astype('O')

# Pure Python
%timeit set(s_all)
451 ms ± 5.94 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) 

%timeit set(s_unicode_obj)
71.9 ms ± 5.91 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

# using set on a list
s_list = s_all.tolist()
%timeit set(s_list)
63.1 ms ± 7.38 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

# NumPy
%timeit np.unique(s_unicode_obj)
1.69 s ± 97.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit np.unique(s_all)
633 ms ± 3.99 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# Pandas
%timeit pd.unique(s_unicode_obj)
97.6 ms ± 6.61 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

# Cython
%timeit unique_cython_str(s_unicode_obj)
60 ms ± 5.81 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

# Cython - c++ unordered_set
%timeit unique_cpp_str2(s_bytes_obj)
247 ms ± 8.45 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

讨论

因此，似乎Python的集合在Unicode字符串方面表现优于pandas，但在整数方面并非如此。另外，在Cython中遍历数组实际上对我们没有帮助。

使用整数欺骗

如果您知道整数范围不太复杂，则可以绕过集合。当遇到每个整数时，您可以简单地创建一个所有零/False的第二个数组，并将它们的位置设置为True并将该数字附加到列表中。由于不进行哈希处理，因此这非常快。

以下内容适用于正整数数组。如果您有负整数，则必须添加一个常量来将数字向上移动到0。

@cython.wraparound(False)
@cython.boundscheck(False)
def unique_bounded(ndarray[np.int64_t] a):
    cdef int i, n = len(a)
    cdef ndarray[np.uint8_t, cast=True] unique = np.zeros(n, dtype=bool)
    cdef list result = []
    for i in range(n):
        if not unique[a[i]]:
            unique[a[i]] = True
            result.append(a[i])
    return result

%timeit unique_bounded(a)
1.18 ms ± 21.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

当然，缺点是内存使用量，因为您的最大整数可能会强制使用非常大的数组。但是，如果您知道每个数字有多少个有效数字，这种方法也可以用于浮点数。

摘要

整数1000000个中有50个独特的

• Pandas-5毫秒
• Python集合列表-10毫秒
• Cython集合-13毫秒
• “作弊”整数-1.2毫秒

字符串1000000个中有1000个独特的

• Cython集合-60毫秒
• Python集合列表-63毫秒
• Pandas-98毫秒

感谢所有帮助加快速度的人。

- Ted Petrou

---

@Dark 是的，Python集合非常快 - 对于字符串比pandas更快，但对于整数则要慢得多。我想知道是否有更高效的方法，也许可以使用Cython并结合C/C++来实现。 - Ted Petrou

Python中的set使用与dict相同的哈希处理方式，这是Python的核心处理组件。而列表迭代比数组迭代更快。因此，我并不惊讶Cython没有改进基本的Python处理过程。它没有为计算带来任何新东西。 - hpaulj

@hpaulj 是的，我正在检查每次迭代调用add设置方法是否最佳。我还想看看是否有其他哈希表实现可以产生更快的结果。而且也想知道“使用整数作弊”的策略是否可行。 - Ted Petrou

请将您的问题中的大部分内容移至答案并标记为已接受，这将成为一个有用的重复目标。 - cs95

Cython 为大多数 C++ 标准库容器定义了自动转换。（显然会有一些复制开销。） - DavidW

Answer 1

我认为回答你的问题“如何快速找到唯一元素”的答案是“这要看情况”。它取决于你的数据集和硬件。

对于你的场景（我主要看了整数情况），pandas（并使用了khash）做得相当不错。我无法使用std :: unordered_map来匹配此性能。

然而，在我的实验中，google :: dense_hash_set比pandas解决方案略快。

请继续阅读以获取更详细的解释。

---

我希望先解释一下您所观察到的结果，并在后面利用这些见解。

我从您的int-example开始：数组中只有50个独特元素，但有100万个元素。

import numpy as np
import pandas as pd
a=np.random.randint(0,50, 10**6, dtype=np.int64)

作为基准，我机器上 np.unique() 和 pd.unique() 的计时如下：

%timeit np.unique(a)
>>>82.3 ms ± 539 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit pd.unique(a)
>>>9.4 ms ± 110 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

pandas的方法使用set(O(n))比numpy的排序方法(O(nlogn))快大约10倍。当n=10**6时，log n = 20，因此这个因素是预期差异。

另一个区别是，np.unique返回一个已排序的数组，因此可以使用二分查找来查找元素。而pd.unique返回一个未排序的数组，所以我们需要将其排序（如果原始数据中没有很多重复项，则可能是O(n log n)），或者将其转换为类似于集合的结构。

让我们来看一下简单的Python-Set：

%timeit set(a)
>>> 257 ms ± 21.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

首先我们必须意识到一个问题：我们正在比较苹果和橙子。之前的unique函数返回由低级c整数组成的numpy数组，而这个函数返回一组完整的Python整数。两者是截然不同的！

这意味着对于numpy数组中的每个元素，我们必须先创建一个Python对象 - 这是一种相当大的开销，只有在此之后才能将其添加到集合中。

将转换为Python整数可以在预处理步骤中完成 - 使用您的版本list：

A=list(a)
%timeit set(A)
>>> 104 ms ± 952 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit set(list(a))
>>> 270 ms ± 23.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

创建Python整数需要100毫秒以上。然而，Python整数比低级的C整数更加复杂，因此处理它们的代价更高。在C整数上使用pd.unique，然后升级为Python集合会更快。

现在是您的Cython版本：

%timeit unique_cython_int(a)
31.3 ms ± 630 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

我其实不是很理解。我希望它的表现类似于set(a) - Cython将消除解释器，但这不能解释10倍的因素。然而，我们只有50个不同的整数（它们在整数池中是偶数，因为它们小于256），所以可能存在一些优化，起到了作用/差异。

让我们尝试另一个数据集（现在有10**5个不同的数字）：

b=np.random.randint(0, 10**5,10**6, dtype=np.int64)
%timeit unique_cython_int(b)
>>> 236 ms ± 31.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit set(b)
>>> 388 ms ± 15.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

速度提升低于2是我所期望的。

让我们看一下cpp版本：

%timeit unique_cpp_int(a)
>>> 25.4 ms ± 534 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit unique_cpp_int(b)
>>> 100 ms ± 4.8 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

这里存在一些开销，需要将数据从cpp-set复制到Python set（正如DavidW所指出的那样），但除此之外，根据我的经验，其行为符合我的预期：std::unordered_map比Python略快，但并不是最好的实现 - panda似乎更胜一筹：

%timeit set(pd.unique(b))
>>> 45.8 ms ± 3.48 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

看起来，在存在许多重复项且哈希函数便宜的情况下，pandas解决方案很难被击败。

或许可以尝试使用谷歌数据结构。

---

然而，当数据只有很少的重复时，numpy的排序解决方案可能会变得更快。主要原因是numpy的unique仅需要两倍的内存 - 原始数据和输出，而pandas的哈希集合解决方案需要更多的内存：原始数据、集合和输出。对于大型数据集，这可能是拥有足够RAM和没有足够RAM之间的区别。

它取决于集合实现需要多少内存开销，这总是在内存和速度之间的权衡。例如，std :: unordered_set 至少需要32字节来保存8字节整数。一些谷歌的数据结构可以做得更好。

使用pandas / numpy unique运行/usr/bin/time -fpeak_used_memory:%M python check_mem.py：

#check_mem.py
import numpy as np
import pandas as pd
c=np.random.randint(0, 2**63,5*10**7, dtype=np.int64)
#pd.unique(c)  
np.unique(c)

显示numpy的大小为1.2 GB，pandas的大小为2.0 GB。

实际上，在我的Windows机器上，如果数组中仅有唯一元素（或接近唯一元素），np.unique比pd.unique更快，即使对于“仅”10^6个元素（可能是因为使用的集合增长时需要重新哈希）。但这在我的Linux机器上并不是这种情况。

---

另一个 pandas 不太适用的场景是计算哈希函数不便宜的情况: 考虑长字符串 (比如说 1000 个字符) 作为对象。

要计算哈希值，需要考虑所有 1000 个字符 (这意味着很多数据 -> 很多哈希缺失)，比较两个字符串通常在一个或两个字符后就完成了 - 那么很有可能我们已经知道这些字符串不同了。因此，numpy 的 unique 的 log n 因子看起来并不那么糟糕。

在这种情况下，使用树集可能会更好。

---

改进cpp-unordered set：

由于其方法reserve()，使用cpp的unordered set方法可以改进，这将消除重新哈希的需要。但是在Cython中使用非常麻烦，因为它没有被导入。

然而，对于仅具有50个唯一元素的数据，保留不会对运行时间产生任何影响，并且对于几乎所有元素都是唯一的数据，最多只会有2倍的因子（由于使用的调整大小策略而导致的摊销成本）。

ints的哈希函数是身份函数（至少对于gcc），因此在这里没有太多可获得的收益（我认为使用更高级的哈希函数也不会有所帮助）。

我看不出如何调整cpp的unordered-set以击败pandas使用的khash实现，后者似乎非常适合此类任务。

这里举例 这些古老的基准测试，显示出 khash 比 std::unordered_map 稍微快一些，只有 google_dense 更快。

---

使用谷歌密集地图：

在我的实验中，谷歌密集地图（此处提供）能够击败khash。基准测试代码可以在答案末尾找到。

如果只有50个唯一元素，效率会更高：

#50 unique elements:
%timeit google_unique(a,r)
1.85 ms ± 8.26 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit pd.unique(a)
3.52 ms ± 33.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

如果只有唯一的元素，那么速度也会更快：

%timeit google_unique(c,r)
54.4 ms ± 375 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
In [3]: %timeit pd.unique(c)
75.4 ms ± 499 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

我的几次实验也表明，google_hash_set可能比khash使用更多的内存（高达20%），但需要进行更多的测试来确定是否真的如此。

---

我不确定我的回答是否对你有所帮助。我的总结如下：

1. 如果我们需要一组Python整数，set(pd.unique(...))似乎是一个很好的起点。
2. 在某些情况下，numpy的排序解决方案可能更好（占用更少的内存，有时哈希计算过于昂贵）。
3. 了解更多关于数据的信息可以用来调整解决方案，通过做出更好的权衡（例如使用更少/更多的内存/预分配，以便我们不需要重新哈希或使用位集进行查找）。
4. Pandas的解决方案似乎已经针对一些常见情况进行了调整，但是对于其他情况，另一种权衡可能更好——google_dense是最有前途的候选方案。

---

谷歌测试的列表：

#google_hash.cpp
#include <cstdint>
#include <functional>
#include <sparsehash/dense_hash_set>

typedef int64_t lli;
void cpp_unique(lli *input, int n, lli *output){

  google::dense_hash_set<lli, std::hash<lli> > set;
  set.set_empty_key(-1);
  for (int i=0;i<n;i++){
     set.insert(input[i]);
  }  

  int cnt=0;
  for(auto x : set)
    output[cnt++]=x;
}

对应的pyx文件：

#google.pyx
cimport numpy as np
cdef extern from "google_hash.cpp":
    void cpp_unique(np.int64_t  *inp, int n, np.int64_t *output)

#out should have enough memory:
def google_unique(np.ndarray[np.int64_t,ndim=1] inp, np.ndarray[np.int64_t,ndim=1] out):
    cpp_unique(&inp[0], len(inp), &out[0])

setup.py文件：

from distutils.core import setup, Extension
from Cython.Build import cythonize
import numpy as np

setup(ext_modules=cythonize(Extension(
            name='google',
            language='c++',
            extra_compile_args=['-std=c++11'],
            sources = ["google.pyx"],
            include_dirs=[np.get_include()]
    )))

在调用python setup.py build_ext --inplace后，执行Ipython-benchmark脚本：

import numpy as np
import pandas as pd
from google import google_unique

a=np.random.randint(0,50,10**6,dtype=np.int64)
b=np.random.randint(0, 10**5,10**6, dtype=np.int64)
c=np.random.randint(0, 2**63,10**6, dtype=np.int64)
r=np.zeros((10**6,), dtype=np.int64)

%timeit google_unique(a,r
%timeit pd.unique(a)

---

其他列表

修复后的Cython版本：

%%cython
cimport cython
from numpy cimport ndarray
from cpython cimport set
cimport numpy as np
@cython.wraparound(False)
@cython.boundscheck(False)
def unique_cython_int(ndarray[np.int64_t] a):
    cdef int i
    cdef int n = len(a)
    cdef set s = set()
    for i in range(n):
        s.add(a[i])
    return s

C++修复后的版本：

%%cython -+ -c=-std=c++11
cimport cython
cimport numpy as np
from numpy cimport ndarray
from libcpp.unordered_set cimport unordered_set
@cython.wraparound(False)
@cython.boundscheck(False)
def unique_cpp_int(ndarray[np.int64_t] a):
    cdef int i
    cdef int n = len(a)
    cdef unordered_set[int] s
    for i in range(n):
        s.insert(a[i])
    return s

- ead

---

非常感谢您提供了如此详细的解释和 Google Dense Hash 的代码。对于整数，我认为我会使用一个简单的策略来映射到布尔数组。此外，Python 的 set 对于字符串似乎比 pandas 更好用。 - Ted Petrou

如果您的整数范围已知且较小，那么可能没有什么比布尔方法更好了。抱歉我没有深入研究字符串情况（它更加复杂）。最终，对于集合实现可能存在不同的权衡，并且取决于数据哪个权衡是“正确”的。然而，了解不同集合实现的优缺点会有所帮助。 - ead