有几个有趣的解决方案不依赖于groupby
. 第一个非常简单:
def apply_to_bins(func, values, bins):
return zip(*((bin, func(values[bins == bin])) for bin in set(bins)))
这使用“花式索引”而不是分组,并且对于小输入执行得相当好;基于列表理解的变体做得更好(请参阅下面的时间)。
def apply_to_bins2(func, values, bins):
bin_names = sorted(set(bins))
return bin_names, [func(values[bins == bin]) for bin in bin_names]
这些具有可读性强的优点。两者的表现也比groupby
小输入好,但对于大输入,它们的速度要慢得多,尤其是当有很多 bin 时;他们的表现是O(n_items * n_bins)
。对于小输入,另一种numpy
基于 - 的方法较慢,但对于大输入要快得多,尤其是对于具有大量 bin 的大输入:
def apply_to_bins3(func, values, bins):
bins_argsort = bins.argsort()
values = values[bins_argsort]
bins = bins[bins_argsort]
group_indices = (bins[1:] != bins[:-1]).nonzero()[0] + 1
groups = numpy.split(values, group_indices)
return numpy.unique(bins), [func(g) for g in groups]
一些测试。首先对于小输入:
>>> def apply_to_bins_groupby(func, x, b):
... return zip(*[(k, np.product(x[list(v)]))
... for k, v in groupby(np.argsort(b), key=lambda i: b[i])])
...
>>> x = numpy.array([1, 2, 3, 4, 5, 6])
>>> b = numpy.array(['a', 'b', 'a', 'a', 'c', 'c'])
>>>
>>> %timeit apply_to_bins(numpy.prod, x, b)
10000 loops, best of 3: 31.9 us per loop
>>> %timeit apply_to_bins2(numpy.prod, x, b)
10000 loops, best of 3: 29.6 us per loop
>>> %timeit apply_to_bins3(numpy.prod, x, b)
10000 loops, best of 3: 122 us per loop
>>> %timeit apply_to_bins_groupby(numpy.prod, x, b)
10000 loops, best of 3: 67.9 us per loop
这里的apply_to_bins3
表现不太好,但它仍然比最快的慢一个数量级。变大时效果更好n_items
:
>>> x = numpy.arange(1, 100000)
>>> b_names = numpy.array(['a', 'b', 'c', 'd'])
>>> b = b_names[numpy.random.random_integers(0, 3, 99999)]
>>>
>>> %timeit apply_to_bins(numpy.prod, x, b)
10 loops, best of 3: 27.8 ms per loop
>>> %timeit apply_to_bins2(numpy.prod, x, b)
10 loops, best of 3: 27 ms per loop
>>> %timeit apply_to_bins3(numpy.prod, x, b)
100 loops, best of 3: 13.7 ms per loop
>>> %timeit apply_to_bins_groupby(numpy.prod, x, b)
10 loops, best of 3: 124 ms per loop
当n_bins
上升时,前两种方法需要很长时间才能在这里显示 - 大约五秒钟。apply_to_bins3
是这里的明显赢家。
>>> x = numpy.arange(1, 100000)
>>> bn_product = product(['a', 'b', 'c', 'd', 'e'], repeat=5)
>>> b_names = numpy.array(list(''.join(s) for s in bn_product))
>>> b = b_names[numpy.random.random_integers(0, len(b_names) - 1, 99999)]
>>>
>>> %timeit apply_to_bins3(numpy.prod, x, b)
10 loops, best of 3: 109 ms per loop
>>> %timeit apply_to_bins_groupby(numpy.prod, x, b)
1 loops, best of 3: 205 ms per loop
总的来说,groupby
在大多数情况下可能没问题,但不太可能很好地扩展,正如这个线程所建议的那样。使用 pure(er)numpy
方法,对于小输入来说速度较慢,但只有一点点;权衡是一个很好的选择。