加速Python循环

I'd use numba to speed up the computation (with added short-circuit: when the two sums are already >= 2 there's no need to count further):

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 5700 x):

0.023107149987481534