如何用向量化方法实现Pandas中生成新列x的指定逻辑？

问题背景 #

我有如下Pandas DataFrame：

import pandas as pd
df = pd.DataFrame(
    {
        'start': [3, 11, 9, 19, 22],
        'end': [10, 17, 10, 25, 30]
    }
)

需要生成新列x，预期输出如下：

start  end    x
0      3   10    10
1     11   17    17
2      9   10    NaN
3     19   25    25
4     22   30    NaN

逻辑说明 #

逐行处理规则：

• 第0行的x为df.end.iloc[0]，该值作为"保留值"
• 后续行若start大于当前保留值，则该行x设为自身end，同时更新保留值为这个end
• 若start小于等于当前保留值，该行x设为NaN

举个例子：

• 初始保留值10，第1行start=11>10，故x=17，保留值更新为17
• 第2行start=9<17，x=NaN
• 保留值17，第3行start=19>17，x=25，保留值更新为25
• 第4行start=22<25，x=NaN

更多示例 #

以下是多组不同DataFrame及对应预期输出：
示例1：

df = pd.DataFrame({'start': [3, 20, 11, 19, 22],'end': [10, 17, 21, 25, 30]})

预期输出：

start  end     x
0      3   10  10.0
1     20   17  17.0
2     11   21   NaN
3     19   25   25.0
4     22   30   NaN

示例2：

df = pd.DataFrame({'start': [3, 9, 11, 19, 22],'end': [10, 17, 21, 25, 30]})

预期输出：

start  end     x
0      3   10   10.0
1      9   17   NaN
2     11   21   21.0
3     19   25   NaN
4     22   30   30.0

示例3：

df = pd.DataFrame({'start': [3, 11, 9, 19, 22],'end': [10, 17, 21, 25, 30]})    

预期输出：

start  end     x
0      3   10  10.0
1     11   17  17.0
2      9   21   NaN
3     19   25  25.0
4     22   30   NaN

当前循环实现 #

我已通过如下循环代码实现该功能：

import numpy as np
l = []
for ind, row in df.iterrows():
    if ind == 0:
        x = row['end']
        l.append(x)
        continue
    if row['start'] > x:
        x = row['end']
        l.append(x)
    else:
        l.append(np.NaN)
df['x'] = l

提问 #

是否存在向量化的实现方式？

向量化实现方案 #

可以通过累积最大值结合布尔掩码实现完全向量化操作，无需循环，大数据量下效率更高：

完整代码 #

import pandas as pd
import numpy as np

df = pd.DataFrame(
    {
        'start': [3, 11, 9, 19, 22],
        'end': [10, 17, 10, 25, 30]
    }
)

# 初始化掩码：标记需要更新保留值的行，第一行默认标记为True
mask = np.zeros(len(df), dtype=bool)
mask[0] = True

# 计算动态累积的保留值（仅在标记行更新）
current_max = df['end'].where(mask).cummax()

# 后续行的标记条件：start > 前一行的累积保留值
mask[1:] = df['start'].iloc[1:] > current_max.iloc[:-1]

# 生成x列：标记行保留end值，其余为NaN
df['x'] = df['end'].where(mask)
print(df)

输出结果：

start  end     x
0      3   10  10.0
1     11   17  17.0
2      9   10   NaN
3     19   25  25.0
4     22   30   NaN

原理说明 #

1. 掩码标记：用布尔数组标记哪些行需要更新保留值（这些行的x就是自身end）
2. 动态累积最大值：通过cummax()结合掩码，实时更新当前的保留值，确保每一行判断时用的是最新的保留值
3. 生成结果列：利用where()方法，仅保留标记行的end值，其余设为NaN

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