使用cython加速itertools组合

时间:2018-05-06 09:53:20

标签: python cython itertools

我有以下代码使用itertools生成指定范围内的所有可能组合,但我无法通过使用cython代码获得任何速度提升。原始代码是这样的:

from itertools import *

def x(e,f,g):
    a=[]        
    for c in combinations(range(e, f),g):
        d = list((c))
        a.append(d) 

并在声明cython的类型之后:

from itertools import *

cpdef x(int e,int f,int g):        
    cpdef tuple c
    cpdef list a
    cpdef list d

    a=[]        
    for c in combinations(range(e, f),g):    
        d = list((c))
        a.append(d)

我将后者保存为test_cy.pyx并使用cythonize -a -i test_cy.pyx

进行编译

编译完成后,我使用以下代码创建了一个新脚本并运行它:

import test_cy

test_cy.x(1,45,6) 

我没有得到任何显着的速度改进,仍然与原始剧本大约相同的时间,大约10.8秒。

有什么我做错了或者是已经如此优化的itertools,以至于它的速度无法有任何更大的改进?

1 个答案:

答案 0 :(得分:2)

正如评论中已经指出的那样,你不应该期望cython加速你的代码,因为算法花费在迭代工具和列表创建的大部分时间。

因为我很高兴看到itertools的通用实现如何对抗老派的伎俩,让我们来看看这个“n中所有子集k”的Cython实现:

%%cython

ctypedef unsigned long long ull

cdef ull next_subset(ull subset):
   cdef ull smallest, ripple, ones
   smallest = subset& -subset      
   ripple = subset + smallest    
   ones = subset ^ ripple    
   ones = (ones >> 2)//smallest 
   subset= ripple | ones    
   return subset


cdef subset2list(ull subset, int offset, int cnt):  
    cdef list lst=[0]*cnt #pre-allocate
    cdef int current=0;
    cdef int index=0
    while subset>0:
        if((subset&1)!=0):
            lst[index]=offset+current
            index+=1
        subset>>=1
        current+=1
    return lst

def all_k_subsets(int start, int end, int k):
    cdef int n=end-start      
    cdef ull MAX=1L<<n;
    cdef ull subset=(1L<<k)-1L;
    lst=[]
    while(MAX>subset):
         lst.append(subset2list(subset, start, k))
         subset=next_subset(subset)
    return lst

这个实现使用了一些众所周知的bit-tricks并且有限制,它只适用于最多64个元素。

如果我们比较两种方法:

>>> %timeit x(1,45,6)
2.52 s ± 108 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

>>> %timeit all_k_subsets(1,45,6)
1.29 s ± 5.16 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

因素2的加速非常令人失望。

然而,瓶颈是列表的创建,而不是计算本身 - 很容易检查,没有列表创建,计算大约需要0.1秒。

我对它的看法:如果你对速度很认真,你不应该创建这么多的列表,而是动态地继续运行子集(最好是在cython中) - 可以加速超过10。如果必须将所有子集都作为列表,那么你不能期望大幅加速。