我有一些代码,它循环处理了大量的 itertools.combinations
现在是一个性能瓶颈。我正试图转向 numba
's @jit(nopython=True)
以加快它的速度,但我遇到了一些问题。
首先,似乎numba不能处理 itertools.combinations
本身,按照这个小例子。
import itertools
import numpy as np
from numba import jit
arr = [1, 2, 3]
c = 2
@jit(nopython=True)
def using_it(arr, c):
return itertools.combinations(arr, c)
for i in using_it(arr, c):
print(i)
throw error: numba.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Unknown attribute 'combinations' of type Module(<module 'itertools' (built-in)>)
经过一些谷歌搜索,我发现 这个github问题 其中提问者提出了这个numba-safe函数来计算组合。
@jit(nopython=True)
def permutations(A, k):
r = [[i for i in range(0)]]
for i in range(k):
r = [[a] + b for a in A for b in r if (a in b)==False]
return r
利用这个函数,我可以很容易地筛选出组合。
@jit(nopython=True)
def combinations(A, k):
return [item for item in permutations(A, k) if sorted(item) == item]
现在我可以运行这个 combinations
函数而不出错,并得到正确的结果。然而,现在用了 @jit(nopython=True)
比没有它。运行这个计时测试。
A = list(range(20)) # numba throws 'cannot determine numba type of range' w/o list
k = 2
start = pd.Timestamp.utcnow()
print(combinations(A, k))
print(f"took {pd.Timestamp.utcnow() - start}")
时钟在2. 6秒的情况下。numba @jit(nopython=True)
装饰者,而且在他们评论出来的情况下,一秒钟的时间不到1000。所以这对我来说也不是一个可行的解决方案。
在这种情况下,使用Numba并没有什么好处,因为 itertools.combinations
是用C写的。
如果你想对它进行基准测试,这里有一个Numba Python的实现,它是用C语言编写的。itertools.combinatiions
的。
@jit(nopython=True)
def using_numba(pool, r):
n = len(pool)
indices = list(range(r))
empty = not(n and (0 < r <= n))
if not empty:
result = [pool[i] for i in indices]
yield result
while not empty:
i = r - 1
while i >= 0 and indices[i] == i + n - r:
i -= 1
if i < 0:
empty = True
else:
indices[i] += 1
for j in range(i+1, r):
indices[j] = indices[j-1] + 1
result = [pool[i] for i in indices]
yield result
在我的机器上,这个速度是15倍左右的速度 itertools.combinations
. 获取排列组合和过滤组合的速度肯定会更慢。