我需要使用 numpy 的矢量化来优化我的双循环答案

【问题标题】：I need to use numpy's vectorization to optimize my double for loop我需要使用 numpy 的矢量化来优化我的双循环
【发布时间】：2021-06-17 00:47:35
【问题描述】：

我有一个带有嵌套 for 循环的 python 函数，它被调用了数千次，而且速度太慢。从我在网上阅读的内容来看，应该有一种方法可以使用 numpy 向量化对其进行优化，以便以更快的 C 代码而不是 python 进行迭代。但是，我以前从未使用过 numpy，我无法弄清楚。

函数是：从第一行开始切片，然后从后到前计算s列第1行的和，如果大于等于n，返回总数；然后从前两行切片，然后从后到前计算s列第1行的和，如果小于n，再计算s列第2行的和，如果大于等于n ，返回行数；以此类推，直到所有行都被循环。

原代码（python 3.8）为：

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

def sumbars(df, s, n):
    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']

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()

我需要优化的双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 向量化优化最后一个 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

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

【问题讨论】：

并非每个计算都是“可向量化的”。 numpy 向量化意味着使用已编译的 numpy 方法。这些方法中的大多数本质上是“并行的”，对整个数组进行操作而没有任何隐含的顺序或顺序（尽管c 代码会迭代）。 np.cumsum 是一个例外，一个本质上是顺序的函数。您对cumsum 的使用表明问题本质上是连续的，并且可能需要跳出框框思考以“矢量化”（如果可能的话）。
我添加了原始代码。除了使用 cumsum，还有其他方法可以优化我的双 for 循环吗？
您可以随时尝试numba 库，它应该可以很好地与 numpy 函数配合使用。
df[s] 的输入值总是正数吗？
df[s] 的输入值总是正数

标签： python pandas performance dataframe numpy

【解决方案1】：

假设 df[s] 的输入值始终为正数，您可以显着改进算法。实际上，您当前实现的复杂度是 O(n**2) 而在这种情况下，可以编写复杂度为 O(n log n) 的算法（甚至可能是 @987654324 @)。

改进的解决方案在于只计算一次cumsum 并执行增量更新。实际上，不需要重新计算它，因为np.cumsum(arr[i:]) 总是等于fullCumsum[i:]-m 与fullCumsum = np.cumsum(arr) 和m = np.sum(arr[:i])。 fullCumsum 独立于 i，因此可以预先计算。 m 可以在循环中增量计算。 k 可以使用二进制搜索找到，因为值已排序。

这是最终的（纯 Python）实现：

def sumbarsFast(df, s, n, ):
    sdata = np.array(df[s])
    sdata = np.flipud(sdata)
    l = len(sdata)
    sumbars = np.zeros(l)
    fullCumsum = np.cumsum(sdata)
    m = 0
    for i in range(l):
        k = np.searchsorted(fullCumsum[i:], n + m)
        if k == len(fullCumsum)-i:
            k = 0
        if (k != 0) or ((fullCumsum[i+k] != m) and (n == 1)):
            sumbars[l - i - 1] = k + 1
        m += sdata[i]
    sumbars = sumbars.astype(int)
    df['sumbars'] = sumbars
    return df['sumbars']

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

from numba import njit, prange, int64

@njit(int64[:](int64[:], int64), parallel=True)
def sumbarsKernel(sdata, n):
    l = len(sdata)
    fullCumsum = np.cumsum(sdata)
    sumbars = np.zeros(l, dtype=np.int64)
    nShared = np.int64(n)
    for i in prange(l):
        m = fullCumsum[i-1] if i > 0 else 0
        k = np.searchsorted(fullCumsum[i:], n + m)
        if k == len(fullCumsum)-i:
            k = 0
        if (k != 0) or ((fullCumsum[i+k] != m) and (nShared == 1)):
            sumbars[l - i - 1] = k + 1
    return sumbars

def sumbarsSuperFast(df, s, n, ):
    sdata = np.array(df[s], dtype=np.int64)
    sdata = np.flipud(sdata)
    sumbars = sumbarsKernel(sdata, n)
    sumbars = sumbars.astype(int)
    df['sumbars'] = sumbars
    return df['sumbars']

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

Your naive version:         589.764 seconds
Your "optimized" version:   121.683 seconds
sumbarsFast (pure Python):    1.081 second
sumbarsSuperFast (Numba):     0.004 second

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

【讨论】：

看来效果超级好，我现在要调试代码。
如果我使用我发布的数据框，我注释掉了 JIT 行并且 sumbarsSuperFast 的结果是正确的。如果没有注释，结果是错误的。这是为什么呢？
确实，问题来自 Numba 代码中的并行性。它看起来像 Numba 中的错误 或 prange 的限制（与 Python 对象有关）...我通过显式复制 Python 整数 n 解决了这个问题到一个安全的共享 np.int64 变量中。现在所有实现之间的结果都相同，性能仍然相同:)。