使用 itertools.groupby 性能进行 NumPy 分组答案

【问题标题】：NumPy grouping using itertools.groupby performance使用 itertools.groupby 性能进行 NumPy 分组
【发布时间】：2011-01-10 21:27:44
【问题描述】：

我有许多包含重复项的大型 (>35,000,000) 整数列表。我需要对列表中的每个整数进行计数。以下代码有效，但似乎很慢。其他人可以使用 Python 和 NumPy 来改进基准测试吗？

def group():
    import numpy as np
    from itertools import groupby
    values = np.array(np.random.randint(0,1<<32, size=35000000), dtype='u4')
    values.sort()
    groups = ((k, len(list(g))) for k,g in groupby(values))
    index = np.fromiter(groups, dtype='u4,u2')

if __name__=='__main__':
    from timeit import Timer
    t = Timer("group()","from __main__ import group")
    print t.timeit(number=1)

返回：

$ python bench.py
111.377498865

根据回复：

def group_original():
    import numpy as np
    from itertools import groupby
    values = np.array(np.random.randint(0, 1<<32, size=35000000), dtype='u4')
    values.sort()
    groups = ((k, len(list(g))) for k,g in groupby(values))
    index = np.fromiter(groups, dtype='u4,u2')

def group_gnibbler():
    import numpy as np
    from itertools import groupby
    values = np.array(np.random.randint(0, 1<<32, size=35000000), dtype='u4')
    values.sort()
    groups = ((k,sum(1 for i in g)) for k,g in groupby(values))
    index = np.fromiter(groups, dtype='u4,u2')

def group_christophe():
    import numpy as np
    values = np.array(np.random.randint(0, 1<<32, size=35000000), dtype='u4')
    values.sort()
    counts=values.searchsorted(values, side='right') - values.searchsorted(values, side='left')
    index = np.zeros(len(values), dtype='u4,u2')
    index['f0'] = values
    index['f1'] = counts
    # Erroneous result!

def group_paul():
    import numpy as np
    values = np.array(np.random.randint(0, 1<<32, size=35000000), dtype='u4')
    values.sort()
    diff = np.concatenate(([1], np.diff(values)))
    idx = np.concatenate((np.where(diff)[0], [len(values)]))
    index = np.empty(len(idx)-1, dtype='u4,u2')
    index['f0'] = values[idx[:-1]]
    index['f1'] = np.diff(idx)

if __name__=='__main__':
    from timeit import Timer
    timings=[
                ("group_original", "Original"),
                ("group_gnibbler", "Gnibbler"),
                ("group_christophe", "Christophe"),
                ("group_paul", "Paul"),
            ]
    for method,title in timings:
        t = Timer("%s()"%method,"from __main__ import %s"%method)
        print "%s: %s secs"%(title, t.timeit(number=1))

返回：

$ python bench.py
Original: 113.385262966 secs
Gnibbler: 71.7464978695 secs
Christophe: 27.1690568924 secs
Paul: 9.06268405914 secs

虽然 Christophe 目前给出的结果不正确。

【问题讨论】：

对可能的整数范围是否有任何限制？可以出现所有 2^32 个可能的整数吗？
group()的输出需要key排序吗？
嗨，Sven，每个 2^32 整数出现的机会均等，并且分组输出（即索引）确实需要按升序排列。 values.sort() 并不是真正的瓶颈，它是 group() 的最后一行，它有点慢！干杯！
如果你只想获取整数的频率计数，np.bincount 正在做。 np.bincount 返回 range(max(value)) 中所有整数的计数，即使是零计数，这可能不是您想要的，但它很快。
您好 user333700。我认为 0 到 2^32 之间的值范围意味着 bincount 将使用比大多数计算机拥有的更多的内存！

标签： python algorithm numpy

【解决方案1】：

我通过以下方式获得了三倍的改进：

def group():
    import numpy as np
    values = np.array(np.random.randint(0, 3298, size=35000000), dtype='u4')
    values.sort()
    dif = np.ones(values.shape, values.dtype)
    dif[1:] = np.diff(values)
    idx = np.where(dif>0)
    vals = values[idx]
    count = np.diff(idx)

【讨论】：

感谢您的回答！运行速度非常快，但 len(vals)==len(count[0]) 最后返回 False，这似乎是一个错误？
是的，这是一个小错误。 count 缺少最终值。 count 应该是：np.diff(inx.tolist+[len(values)]) 或者不那么难看的东西。
可以通过使用idx = np.concatenate((np.where(dif)[0],[len(values)])) 和vals = values[idx[:-1]] 分别代替倒数第三行和倒数第二行来修复错误。这确实是仅使用 numpy 的最佳答案。如果你想要更快，我会推荐使用 Cython。与这个 numpy 代码相比，这很容易做到，并且在速度和内存方面有相当显着的改进。
@bpowah 和 @Justin 已将您的修复程序合并到原始帖子中，此解决方案比原始代码快 11 倍！！！
@Justin 你有可以分享的 cython 示例吗？

【解决方案2】：

自保罗的回答被接受以来已经过去了 5 年多。有趣的是， sort() 仍然是公认解决方案的瓶颈。

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
     3                                           @profile
     4                                           def group_paul():
     5         1        99040  99040.0      2.4      import numpy as np
     6         1       305651 305651.0      7.4      values = np.array(np.random.randint(0, 2**32,size=35000000),dtype='u4')
     7         1      2928204 2928204.0    71.3      values.sort()
     8         1        78268  78268.0      1.9      diff = np.concatenate(([1],np.diff(values)))
     9         1       215774 215774.0      5.3      idx = np.concatenate((np.where(diff)[0],[len(values)]))
    10         1           95     95.0      0.0      index = np.empty(len(idx)-1,dtype='u4,u2')
    11         1       386673 386673.0      9.4      index['f0'] = values[idx[:-1]]
    12         1        91492  91492.0      2.2      index['f1'] = np.diff(idx)

接受的解决方案在我的机器上运行 4.0 秒，使用基数排序降至 1.7 秒。

仅通过切换到基数排序，我的整体速度提高了 2.35 倍。在这种情况下，基数排序比快速排序快 4 倍以上。

请参阅How to sort an array of integers faster than quicksort?，这是由您的问题引起的。

对于我使用line_profiler and kernprof 的分析（@profile 来自那里）。

【讨论】：

听起来很像 numpy 需要一个基数排序补丁 :-)
@Donny 是的。其实已经有一个未解决的问题：github.com/numpy/numpy/issues/6050我会把我的发现贴在那里，很容易抓住这个基数排序并将其放入 numpy。
@profile 装饰器使用什么包？看起来很不错。
@danijar 是的，这很酷。 :) 它是line_profiler and kernprof。

【解决方案3】：

根据要求，这是一个 Cython 版本。我做了两次通过数组。第一个找出有多少唯一元素，以便我的数组可以获取唯一值和适当大小的计数。

import numpy as np
cimport numpy as np
cimport cython

@cython.boundscheck(False)
def dogroup():
    cdef unsigned long tot = 1
    cdef np.ndarray[np.uint32_t, ndim=1] values = np.array(np.random.randint(35000000,size=35000000),dtype=np.uint32)
    cdef unsigned long i, ind, lastval
    values.sort()
    for i in xrange(1,len(values)):
        if values[i] != values[i-1]:
            tot += 1
    cdef np.ndarray[np.uint32_t, ndim=1] vals = np.empty(tot,dtype=np.uint32)
    cdef np.ndarray[np.uint32_t, ndim=1] count = np.empty(tot,dtype=np.uint32)
    vals[0] = values[0]
    ind = 1
    lastval = 0
    for i in xrange(1,len(values)):
        if values[i] != values[i-1]:
            vals[ind] = values[i]
            count[ind-1] = i - lastval
            lastval = i
            ind += 1
    count[ind-1] = len(values) - lastval

到目前为止，排序实际上花费的时间最多。使用我的代码中给出的值数组，排序需要 4.75 秒，唯一值和计数的实际查找需要 0.67 秒。使用 Paul 的代码（但具有相同形式的值数组）的纯 Numpy 代码以及我在评论中建议的修复，查找唯一值和计数需要 1.9 秒（当然，排序仍然需要相同的时间）。

排序占用大部分时间是有意义的，因为它是 O(N log N) 并且计数是 O(N)。您可以比 Numpy 加快排序速度（如果我没记错的话，它使用 C 的 qsort），但是您必须真正知道自己在做什么，这可能不值得。此外，可能有一些方法可以进一步加快我的 Cython 代码速度，但这可能不值得。

【讨论】：

【解决方案4】：

我想最明显但仍未提及的方法是简单地使用collections.Counter。它不是用 groupby 构建大量临时使用的列表，而是对整数进行递增计数。这是一个单线器和 2 倍加速，但仍然比纯 numpy 解决方案慢。

def group():
    import sys
    import numpy as np
    from collections import Counter
    values = np.array(np.random.randint(0,sys.maxint,size=35000000),dtype='u4')
    c = Counter(values)

if __name__=='__main__':
    from timeit import Timer
    t = Timer("group()","from __main__ import group")
    print t.timeit(number=1)

与最初的解决方案相比，我的机器从 136 秒加速到 62 秒。

【讨论】：

【解决方案5】：

这是一个 numpy 解决方案：

def group():
    import numpy as np
    values = np.array(np.random.randint(0,1<<32,size=35000000),dtype='u4')

    # we sort in place
    values.sort()

    # when sorted the number of occurences for a unique element is the index of 
    # the first occurence when searching from the right - the index of the first
    # occurence when searching from the left.
    #
    # np.dstack() is the numpy equivalent to Python's zip()

    l = np.dstack((values, values.searchsorted(values, side='right') - \
                   values.searchsorted(values, side='left')))

    index = np.fromiter(l, dtype='u4,u2')

if __name__=='__main__':
    from timeit import Timer
    t = Timer("group()","from __main__ import group")
    print t.timeit(number=1)

在我的机器上运行大约 25 秒，而您的初始解决方案大约需要 96 秒（这是一个很好的改进）。

可能还有改进的余地，我不经常使用 numpy。

编辑：在代码中添加了一些 cmets。

【讨论】：

嗨克里斯托夫！我认为右左搜索减法逻辑很棒！在原始问题中进行了基准测试，并将输出按摩到结构化数组中以节省内存。
@Donny，我不确定这是否是您想要的。例如，如果 values 等于 [2,2,3]，那么您将得到 array([[[2,2],[2,2],[3,1]]]) 对应 l 和 array([(2L, 2)], dtype=[('f0','<u4'),('f1','<u2')]) 对应 index。您想首先使用类似 np.unique 的东西找到唯一值，然后只使用这些唯一值进行二进制搜索。
@Justin 你是对的，这种方法的结果是错误的！但是，如果我在 np.unique(values) 上进行搜索排序，那么你会得到 [1,1,1,1,1,1,.....] 这也是错误的！
@Donny，这些实际上可能是正确的结果（[1,1,1,1,1,1...]），因为您从 0 到 4294967296 范围内选择了 35000000 个值。您的范围比您选择的值的数量大 122 倍以上。在这种情况下，您获得重复整数的可能性并不是特别高。
@Justin, np.unique(values) 将始终返回 ([1,1,1,1,1,....] 因为它只返回一个值的单个实例。请提交一些代码在一个单独的答案中，我们试一试:)

【解决方案6】：

将 len(list(g)) 替换为 sum(1 for i in g) 可实现 2 倍加速

【讨论】：

感谢 gnibbler，已在原帖中对这一优化进行了基准测试。

【解决方案7】：

这是一个相当老的线程，但我想我会提到对当前接受的解决方案有一个小的改进：

def group_by_edge():
    import numpy as np
    values = np.array(np.random.randint(0,1<<32,size=35000000),dtype='u4')
    values.sort()
    edges = (values[1:] != values[:-1]).nonzero()[0] - 1
    idx = np.concatenate(([0], edges, [len(values)]))
    index = np.empty(len(idx) - 1, dtype= 'u4, u2')
    index['f0'] = values[idx[:-1]]
    index['f1'] = np.diff(idx)

这在我的机器上测试大约快了半秒；不是很大的改进，但值得。此外，我认为这里发生的事情更清楚了。乍一看，两步diff 方法有点不透明。

【讨论】：

【解决方案8】：

在最新版本的 numpy 中，我们有这个。

import numpy as np
frequency = np.unique(values, return_counts=True)

【讨论】：

这符合要求，但我也想找到一种方法来做更一般的 groupby..
@javadba 你的一般 groupby 是什么？有什么例子吗？顺便说一句，通用解决方案总是会牺牲一些时间。
general groupby 保留每个组中的所有行，而不是应用聚合并因此丢失行。我有一个实现并正在调试它。

【解决方案9】：

排序是 theta(NlogN)，我会选择 Python 的哈希表实现提供的摊销 O(N)。只需使用defaultdict(int) 来保持整数的计数，然后对数组进行一次迭代：

counts = collections.defaultdict(int)
for v in values:
    counts[v] += 1

这理论上更快，不幸的是我现在无法检查。分配额外的内存可能会使其实际上比您的解决方案慢。

编辑：如果您需要节省内存，请尝试基数排序，它在整数上比快速排序要快得多（我相信这是 numpy 使用的）。

【讨论】：

您好 Rafal，感谢您的回答！不幸的是，dicts 的空间效率比 numpy 数组低得多，并且由于有 35,000,000 个值，我很快就会在 4GB 笔记本电脑上耗尽内存。尽可能避免分而治之，一次性完成全部批次。
嗨 Rafal，我认为排序速度不是瓶颈，而是计数方法可以输入 35,000,000 个整数。
@Donny 当您查看上面的选项时，您会发现排序实际上需要 5 倍的时间，然后是 diff 的有效计数。所以排序是的瓶颈，而这个基数搜索可能是对已经建议的解决方案的一个很好的扩展。

【解决方案10】：

您可以尝试以下（ab）使用scipy.sparse：

from scipy import sparse
def sparse_bincount(values):
    M = sparse.csr_matrix((np.ones(len(values)), values.astype(int), [0, len(values)]))
    M.sum_duplicates()
    index = np.empty(len(M.indices),dtype='u4,u2')
    index['f0'] = M.indices
    index['f1']= M.data
    return index

这比获胜的答案慢，可能是因为scipy 目前不支持无符号作为索引类型...

【讨论】：