有几个有趣的解决方案不依赖于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 上升时,前两种方法需要很长时间才能在此处显示 - 大约 5 秒。 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 在大多数情况下可能很好,但不太可能像this thread 所建议的那样很好地扩展。使用纯(er)numpy 方法,对于小输入来说速度较慢,但只有一点;权衡是一个很好的选择。