Python 我需要使用numpy'；s矢量化以优化我的双for循环

我有一个带有嵌套for循环的python函数，它被调用了数千次，速度太慢了。从我在网上读到的内容来看，应该有一种方法可以通过numpy矢量化对其进行优化，这样迭代就可以用比python快得多的C代码完成。但是，我以前从来没有和numpy一起工作过，我想不出来

函数为：先从第一行开始切片，然后从后向前计算s列1行的和，如果大于或等于n，则返回总数；然后从前两行开始切片，然后从后向前计算s列1行之和，如果小于n，则计算s列2行之和，如果大于或等于n，则返回行数；依此类推，直到所有行都循环

原始代码（python 3.8）是：

我需要优化双for循环：

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

有人帮我优化了我的代码一次，我需要优化最后一个for循环一次。 优化for循环后：

import pandas as pd
from time import time
import numpy as np


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):
        n1 = sdata[i:]
        cumsum = np.cumsum(n1)
        k = np.argmax(cumsum >= n)
        n2 = cumsum[k]
        if (k != 0) | ((n2 != 0) & (n == 1)):
            sumbars[l - i - 1] = k + 1
    sumbars = sumbars.astype(int)
    df['sumbars'] = sumbars
    stop = time()
    global sumbarstime
    sumbarstime += stop - start
    # df.to_csv("my.csv")
    # print(df['sumbars'])
    return df['sumbars']

def main():
    # df=pd.read_csv('todo.csv')
    df = pd.DataFrame({'ts_code':['IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX','IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX','IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX','IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX','IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX','IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX','IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX','IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX','IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX','IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX','IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX','IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX','IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX','IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX','IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX','IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX', 'IC2103.CFX'],'jc':[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0,0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1,0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0,1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1]})
    df['jcts2'] = sumbars(df,'jc',2)

    
    print('')
    df.to_csv("result.csv",index=False,sep=',')

    # print sec

    return

main()

帮助：我需要使用numpy矢量化优化last for循环。除了使用cumsum，还有其他方法优化我原来的double for循环吗

    for i in range(l):
        n1 = sdata[i:]
        cumsum = np.cumsum(n1)
        k = np.argmax(cumsum >= n)
        n2 = cumsum[k]
        if (k != 0) | ((n2 != 0) & (n == 1)):
            sumbars[l - i - 1] = k + 1

---

在指出我的问题后，请给我更多的时间来回答。

假设

df[s]

的输入值始终为正，您可以显著改进算法。实际上，当前实现的复杂性是

O（n**2）

，而在这种情况下，可以编写复杂性为

O（n log n）

（甚至可能是

O（n）

）的算法

改进的解决方案包括只计算一次

cumsum

，并执行增量更新。 实际上，不需要重新计算它，因为

np.cumsum（arr[i:]）

始终等于

fullCumsum[i:]-m

，其中

fullCumsum=np.cumsum（arr）

和

m=np.sum（arr[:i]）

。

fullCumsum

独立于

i

，因此可以对其进行预计算。

m

可以在循环中增量计算。

k

可以使用二进制搜索找到，因为值是经过排序的

下面是结果（纯Python）实现：

def sumbarsFast（方位角、方位角、方位角）：
sdata=np.数组（df[s]）
sdata=np.flipud（sdata）
l=len（sdata）
sumbars=np.零（l）
fullCumsum=np.cumsum（sdata）
m=0
对于范围（l）中的i：
k=np.searchsorted（fullCumsum[i:]，n+m）
如果k==len（fullCumsum）-i：
k=0
如果（k！=0）或（fullCumsum[i+k]！=m）和（n==1））：
萨姆巴[l-i-1]=k+1
m+=sdata[i]
sumbars=sumbars.astype（int）
df['sumbars']=sumbars
返回df['sumbars']

使用Numba JIT和并行性可以更快：

来自numba import njit、prange、int64
@njit（int64[：]（int64[：]，int64），parallel=True）
def sumbarsKernel（美国国家数据局，n）：
l=len（sdata）
fullCumsum=np.cumsum（sdata）
sumbars=np.zero（l，dtype=np.int64）
n共享=np.int64（n）
对于prange中的i（l）：
如果i>0，则m=完整总和[i-1]，否则为0
k=np.searchsorted（fullCumsum[i:]，n+m）
如果k==len（fullCumsum）-i：
k=0
如果（k！=0）或（fullCumsum[i+k]！=m）和（nShared==1））：
萨姆巴[l-i-1]=k+1
回流水坑
def SUMBARSSUPPERFAST（东风、南风、北风）：
sdata=np.array（df[s]，dtype=np.int64）
sdata=np.flipud（sdata）
sumbars=sumbarsKernel（美国国家数据局，n）
sumbars=sumbars.astype（int）
df['sumbars']=sumbars
返回df['sumbars']

以下是我的机器上包含200000行的数据帧的计时：

你的天真版本：589.764秒
您的“优化”版本：121.683秒
sumbarsFast（纯Python）：1.081秒
SumbarsuperFast（Numba）：0.004秒

---

最终实现比原始版本快约150000倍，比“优化”版本快约30000倍。

并非所有计算都是“矢量化”的

numpy

矢量化意味着使用编译的numpy方法。大多数方法本质上是“并行”的，在整个数组上运行，没有任何隐含的顺序或序列（尽管

c

代码将迭代）

np.cumsum

是一个例外，本质上是一个顺序函数。使用

cumsum

表明问题本质上是顺序性的，可能需要跳出框框进行“矢量化”（如果可能的话）。我添加了原始代码。除了使用cumsum，还有其他方法优化我的double for loop吗？你可以试试这个库，它应该可以很好地处理numpy函数。

df[s]

的输入值总是正的吗？df[s]的输入值总是正的似乎效果非常好，我现在要调试代码。如果我使用我发布的数据帧，我注释掉了JIT行，sumbarsupperfast的结果是正确的。如果没有注释，则结果是错误的。这是为什么呢？事实上，问题来自Numba代码中的并行性。它看起来像是Numba中的一个bug或是

prange

（与Python对象相关）的一个限制。。。我通过将Python整数

n

显式复制到一个安全的共享

np.int64

变量中，解决了这个问题。现在，所有实现之间的结果都相同，并且性能仍然相同：）。

    for i in range(l):
        n1 = sdata[i:]
        cumsum = np.cumsum(n1)
        k = np.argmax(cumsum >= n)
        n2 = cumsum[k]
        if (k != 0) | ((n2 != 0) & (n == 1)):
            sumbars[l - i - 1] = k + 1