pandas 向量化优化双重for循环

Question

我用pandas 写了个sumbars，二次for循环，感觉效率很慢，需要用向量化优化一下，麻烦各位看看

 

def sumbars(df, s, n):
#s表示列名
    start = time()
    df2 = df[s].copy()
    # df2.loc[0]=1000
    df['sumbars'] = 0
    for i in range(len(df)):
        df3 = df2[:i + 1]
        l = len(df3)
        for j in range(l):
            # if i>=8:
            #     print()
            df4 = df3[l - j - 1:]
            if df4.sum() >= n:
                df.loc[i, 'sumbars'] = j + 1
                # df5=df[['trade_date',s,'sumbars']]
                break

    stop = time()
    global sumbarstime
    sumbarstime = sumbarstime + stop - start
    return df['sumbars']

Answer 1

题主，你的意思是希望找到第s列，以第i行为结尾，满足第j行开始，和大于n的那一个j吗？

如果是这样的话，算法运行效率比较慢的原因是，在if df4.sum() >= n的时候，每次计算完sum(0:i+1)之后，又重新计算一次sum(1:i+1)，再计算sum(2:i+1)，再计算sum(3:i+1)，没有充分利用上一次得到的计算结果。这样的运算之后需要计算1+2+3+...+i+1,是O(n方)级别的。在加上外层的循环，总的时间复杂度是O(n^3).

你可以计算第一次sum(0:i+1)之后，每次减去a[0],a[1],a[2]。。。这样总的时间复杂度就降低为O(n^2)级别的了。

def sumbars(df,s,n):
    #s表示列名 
    start = time()
    sdata = np.array(df[s])
    sdata = np.flipud(sdata)
    l = len(sdata)
    sumbars = np.zeros(l)
    for i in range(l):
        cumsum = np.cumsum(sdata[i:])
        k = np.argmax(cumsum>=n)
        if k != 0:
            sumbars[l-i-1] = k+1
    stop = time()
    print(stop-start)
    global sumbarstime
    sumbarstime = sumbarstime + stop - start
    sumbars.astype(int)
    df['sumbars'] = sumbars
    # df.to_csv("my.csv")
    # print(df['sumbars'])
    return df['sumbars']

更进一步，这个allSum也可以不用每个i都计算一次。只需要每次增加就行。 

def sumbars(df, s, n):
    #s表示列名
    start = time()
    df['sumbars'] = 0
    allSum = 0
    for i in range(len(df)):
        allSum += df[s].iloc[i]
        eachSum = allSum
        for j in range(i+1):
            if eachSum >= n:
                df.loc[i, 'sumbars'] = j + 1
                break
            eachSum -= df[s].iloc[j]
    stop = time()
    global sumbarstime
    sumbarstime = sumbarstime + stop - start
    return df['sumbars']

另外，导致题主之前代码运行比较慢的原因，也可能是多次对df进行拷贝导致的。当DataFrame较大的时候，拷贝也会占用一定时间。

如果我对题主代码的意思理解有偏差，也请评论或私信告诉我。

 

========

和题主沟通后，题主的意思是找到最后一个符合大于等于n的行下标，同时使用numpy进行优化。优化后的代码如下。

def sumbars(df, s, n):
    start = time()
    sdata = np.array(df[s])
    sdata = np.flipud(sdata)
    l = len(sdata)
    sumbars = np.zeros(l)
    for i in range(l):
        cumsum = np.cumsum(sdata[i:])
        k = np.argmax(cumsum>=n)
        if k != 0:
            sumbars[l-i-1] = k+1
    stop = time()
    print(stop-start)
    global sumbarstime
    sumbarstime = sumbarstime + stop - start
    sumbars.astype(int)
    df['sumbars'] = sumbars
    # df.to_csv("my.csv")
    # print(df['sumbars'])
    return df['sumbars']