python - 如何在 Python 中向量化这个 for 循环？

Question

我的数据有 4 列 A-D，其中包含整数。我正在添加一个新列 E，它的第一个值与 D 列中的第一个值相同。如果 E 列中的前一个值为负，则 E 中的下一个值应该是 D 列中的相应值，否则它在列中取相应值C。

import pandas as pd
import numpy as np
from pandas import Series, DataFrame
data=pd.read_excel('/Users/xxxx/Documents/PY Notebooks/Data/yyyy.xlsx')
data1=data.copy()
data1['E']=np.nan
data1.at[0,'E']=data1['D'][0]
l=len(data1)
for i in range(l-1):
    if data1['E'][i]<0:
        data1.at[i+1,'E']=data1['D'][i+1]
    else:
        data1.at[i+1,'E']=data1['C'][i+1]

Answer 1

TL;博士 :转到基准代码并使用方法1。
简答
否 .矢量化是不可能的。
长答案
定理 :对于此特定任务，无法使用任何有限长度的后向滚动窗口来确定给定行的输出，该长度小于到该行的部分长度。
因此，无法以矢量化方式处理此输出逻辑。 (请参阅 this answer 了解在 CPU 中执行矢量化的想法)。输出只能从数据帧的开头计算。
证明 :考虑数据帧的目标行 df .假设有一个大小为 n < partial length 的向后滚动窗口，所以以前的值 df["E"]存在于窗口之前。我们用 state 表示这个先前的值.
考虑一个特殊情况:df["C"] == -1和 df["D"] == 1窗内。

情况 1 ( state < 0 ):此滚动窗口内的输出将为 [1, -1, 1, -1, .....]，使最后一个元素 (-1)^(n-1)

情况 2 ( state >= 0 ):输出将为 [-1, 1, -1, 1, .....]，使最后一个元素 (-1)^(n)

因此，输出 df["E"] 是可能的目标行依赖于窗口外的状态变量。 QED。
有用的答案
虽然矢量化是不可能的，但这并不意味着 显着加速无法实现。一个简单但非常有效的方法是使用 numba-compiled generator执行顺序生成。它只需要将您的逻辑重新写入生成器函数并添加两行:

import numba

@numba.njit
def my_generator_func():
    ....

当然，你可能要 install numba first .如果这是不可能的，那么使用没有 numba 优化的普通生成器也可以。
基准
基准测试在 i5-8250U (4C8T) 笔记本电脑上执行，16GB RAM 运行 64 位 debian 10。Python 版本为 3.7.9，pandas 为 1.1.3。 n = 10^7生成(1000 万)条记录用于基准测试。
结果 :

1. numba-njit: 2.48s
2. plain generator (no numba): 5.13s
3. original: 271.15s

> 100x可以针对原始代码实现效率增益。
代码

from datetime import datetime
import pandas as pd
import numpy as np

n = 10000000  # a large number of rows
df = pd.DataFrame({"C": -np.ones(n), "D": np.ones(n)})
#print(df.head())

# ========== Method 1. generator + numba njit ==========
ti = datetime.now()

import numba

@numba.njit
def gen(plus: np.array, minus: np.array):
    l = len(plus)
    assert len(minus) == l
    # first
    state = minus[0]
    yield state
    # second to last
    for i in range(l-1):
        state = minus[i+1] if state < 0 else plus[i+1]
        yield state

df["E"] = [i for i in gen(df["C"].values, df["D"].values)]

tf = datetime.now()
print(f"1. numba-njit: {(tf-ti).total_seconds():.2f}s")  # 1. numba-njit: 0.47s

# ========== Method 2. Generator without numba ==========
df = pd.DataFrame({"C": -np.ones(n), "D": np.ones(n)})
ti = datetime.now()

def gen_plain(plus: np.array, minus: np.array):
    l = len(plus)
    assert len(minus) == l
    # first
    state = minus[0]
    yield state
    # second to last
    for i in range(l-1):
        state = minus[i+1] if state < 0 else plus[i+1]
        yield state

df["E"] = [i for i in gen_plain(df["C"].values, df["D"].values)]

tf = datetime.now()
print(f"2. plain generator (no numba): {(tf-ti).total_seconds():.2f}s")  #

# ========== Method 3. Direct iteration ==========
df = pd.DataFrame({"C": -np.ones(n), "D": np.ones(n)})
ti = datetime.now()

# code provided by the OP
df['E']=np.nan
df.at[0,'E'] = df['D'][0]
l=len(df)
for i in range(l - 1):
    if df['E'][i] < 0:
        df.at[i+1,'E'] = df['D'][i+1]
    else:
        df.at[i+1,'E'] = df['C'][i+1]

tf = datetime.now()
print(f"3. original: {(tf-ti).total_seconds():.2f}s") # 2. 26.61s