加速Python循环

Question

我创建了一个 for 循环，它从一列中获取一个值，并向前查看该值是否在以后的数据中被超过一两次。该代码可以工作，但由于它运行的数据集非常大，因此代码变得非常慢。我特别怀疑，因为每次迭代都会计算超出某个值的次数（超过大约 500k 行）。有没有办法加快速度？

import pandas as pd

df1 = pd.DataFrame({'index': [0,1,2,3,4], 'Time': ['2022-01-01','2022-01-02','2022-01-03','2022-01-04','2022-01-05'], 'A':[234,456,323,576,234], 'B': [0,1,0,1,0], 'B.v': [0,234,0,323,0], 'in' : [0,0,0,0,0], 'out':[0,0,0,0,0]})


def calc(df1):

    df2 = pd.DataFrame(df1[df1['B'] ==  1])

    for x in range(len(df2)):
        index = df2.iloc[x, df2.columns.get_loc('index')]
        tvalue = df2.iloc[x, df2.columns.get_loc('A')]
        pointvalue = df2.iloc[x, df2.columns.get_loc('B.v')]
        postrates = df1['A'].values[range(index,len(df1))]

        if sum(pointvalue > postrates) == 1:
            df1.iloc[index, df1.columns.get_loc('in')] = 1
        if sum(pointvalue > postrates) >= 2:
            df1.iloc[index, df1.columns.get_loc('in')] = 2

        if sum(tvalue < postrates) == 1:
            df1.iloc[index, df1.columns.get_loc('out')] = 1
        if sum(tvalue < postrates) >= 2:
            df1.iloc[index, df1.columns.get_loc('out')] = 2
    return df1

if __name__ == "__main__":
    print(calc(df1))

Answer 1

我会使用 numba 来加快计算速度（添加 short- Circuit：当两个总和已经 >= 2 时，无需进一步计数）：

import numba


@numba.njit(parallel=True)
def calc_in_out(A, B, Bv, out_in, out_out):
    for idx in numba.prange(len(B)):
        val_b = B[idx]

        if val_b != 1:
            continue

        val_a = A[idx]
        val_Bv = Bv[idx]

        s1, s2 = 0, 0
        for idx2 in range(idx, len(B)):
            s1 += val_Bv > A[idx2]
            s2 += val_a < A[idx2]

            # no need to count further:
            if s1 >= 2 and s2 >= 2:
                break

        if s1 == 1:
            out_in[idx] = 1
        elif s1 >= 2:
            out_in[idx] = 2

        if s2 == 1:
            out_out[idx] = 1
        elif s2 >= 2:
            out_out[idx] = 2

---

基准：

from timeit import timeit

import numba
import numpy as np


@numba.njit(parallel=True)
def calc_in_out(A, B, Bv, out_in, out_out):
    for idx in numba.prange(len(B)):
        val_b = B[idx]

        if val_b != 1:
            continue

        val_a = A[idx]
        val_Bv = Bv[idx]

        s1, s2 = 0, 0
        for idx2 in range(idx, len(B)):
            s1 += val_Bv > A[idx2]
            s2 += val_a < A[idx2]

            # no need to count further:
            if s1 >= 2 and s2 >= 2:
                break

        if s1 == 1:
            out_in[idx] = 1
        elif s1 >= 2:
            out_in[idx] = 2

        if s2 == 1:
            out_out[idx] = 1
        elif s2 >= 2:
            out_out[idx] = 2

def setup_df(N=500_000):
    return pd.DataFrame(
        {
            "Time": ["2022-01-01"] * N,
            "A": np.random.randint(10, 1000, size=N),
            "B": np.random.randint(0, 2, size=N),
            "B.v": np.random.randint(10, 1000, size=N),
            "in": [0] * N,
            "out": [0] * N,
        }
    )


def main():
    df1 = pd.DataFrame(
        {
            "index": [0, 1, 2, 3, 4],
            "Time": [
                "2022-01-01",
                "2022-01-02",
                "2022-01-03",
                "2022-01-04",
                "2022-01-05",
            ],
            "A": [234, 456, 323, 576, 234],
            "B": [0, 1, 0, 1, 0],
            "B.v": [0, 234, 0, 323, 0],
            "in": [0, 0, 0, 0, 0],
            "out": [0, 0, 0, 0, 0],
        }
    )

    # this will compile calc_in_out
    calc_in_out(
        df1["A"].values,
        df1["B"].values,
        df1["B.v"].values,
        df1["in"].values,
        df1["out"].values,
    )
    print(df1)

    to_run = """calc_in_out(
        df1["A"].values,
        df1["B"].values,
        df1["B.v"].values,
        df1["in"].values,
        df1["out"].values,
    )"""

    t = timeit(to_run, setup="df1=setup_df()", globals=globals(), number=1)
    print(t)


if __name__ == "__main__":
    main()

这在我的机器上打印（AMD 5700x）：

0.023107149987481534