numpy.where是否并行修改值？如何无循环实现DataFrame序列更新？

我有如下pandas DataFrame：

value           some other cond
0              0                      true
1              1                      true
2              0                      true
3              1                      true
4              1                      true
5              0                      false

需要对value列执行以下修改：如果前一行的value为0 且当前行的some other cond为true，则将当前行的value改为0，否则保持原值。

我首先尝试用np.where做向量化处理，代码如下：

import pandas as pd
import numpy as np

df["value"] = np.where((df["value"].shift(1) == 0) & (df["Some Other Cond"] == "true"),
    0, df["value"])

得到的结果如下：

value          new value
0              0                0
1              1                -> 0 ( Correct)
2              0                0
3              1                -> 0 ( Correct)
4              1                -> 1 ( Unchanged! even though it should change to 0)
5              0                0

显然numpy.where是并行处理所有单元格，不会逐行从上到下基于更新后的数据处理，不符合预期。

于是我改用循环遍历实现需求，代码如下：

import pandas as pd
import numpy as np

for i in range(1, len(df.index)):
    if (df.at[i - 1, 'value'] == 0) and (df.at[i, 'Some other cond'] == 'true'):
        df.at[i, 'value'] = 0

这个方法能解决问题，但速度太慢：添加循环后，代码运行时间从10秒涨到2分钟，耗时提升了近12倍。

我的问题是：

1. 有没有办法不使用for循环实现该需求？
2. 如果无法避免循环，如何优化循环以提升速度？

补充说明：我采用了Tim Roberts的方案，使用np.apply将运行时间从120秒缩短至35秒，速度提升约4倍。

1. 无循环的向量化实现 #

这个需求属于状态依赖的逐行更新，常规向量化方法无法满足依赖前一行更新后状态的要求，以下两种方法可以解决：

方法一：生成器 + np.fromiter #

用生成器模拟逐行处理逻辑，再通过np.fromiter转换为数组，避免DataFrame索引开销：

import pandas as pd
import numpy as np

def update_values(values, conds):
    prev_val = values[0]
    yield prev_val
    for val, cond in zip(values[1:], conds[1:]):
        if prev_val == 0 and cond == 'true':
            new_val = 0
        else:
            new_val = val
        yield new_val
        prev_val = new_val

# 直接操作numpy数组提升效率
df['value'] = np.fromiter(
    update_values(df['value'].values, df['some other cond'].values),
    dtype=int
)

方法二：numba JIT编译 #

通过numba将循环编译为机器码，速度接近纯C级别，适合超大规模数据集：

import pandas as pd
import numpy as np
from numba import jit

@jit(nopython=True)
def update_numba(values, conds):
    n = len(values)
    result = np.empty(n, dtype=np.int64)
    result[0] = values[0]
    for i in range(1, n):
        # conds为提前转换好的布尔数组
        if result[i-1] == 0 and conds[i]:
            result[i] = 0
        else:
            result[i] = values[i]
    return result

# 先将条件列转为布尔数组适配numba
df['cond_bool'] = df['some other cond'] == 'true'
df['value'] = update_numba(df['value'].values, df['cond_bool'].values)

2. 循环的优化方案 #

如果必须保留循环，核心是减少DataFrame的索引开销，直接操作numpy数组：

import pandas as pd
import numpy as np

# 提取列转为numpy数组，避免循环中频繁访问DataFrame
values = df['value'].values.copy()
conds = df['some other cond'].values

for i in range(1, len(values)):
    if values[i-1] == 0 and conds[i] == 'true':
        values[i] = 0

# 将结果赋值回DataFrame
df['value'] = values

优化点说明：

• 提前把列转为numpy数组，避免循环中df.at的频繁Python-C交互开销
• 用copy()确保原DataFrame的初始值不被修改（若无需保留初始值可省略）

内容的提问来源于stack exchange，提问作者Mohsen Shahhosseini