如何在Pandas中定位掩码首次匹配行并标记其下一行指定列

原始DataFrame #

import pandas as pd 
df = pd.DataFrame(
    {
        'a': [100, 1123, 123, 100, 1, 0, 1],
        'b': [1000, 11123, 1123, 0, 55, 0, 1],
        'c': ['a', 'b', 'c', 'd', 'e', 'f', 'g'],
    }
)

规则说明 #

• 生成掩码 mask = (df.a > df.b)，找到首次满足该掩码的行
• 将该行下一行的c列值，赋值给该行的x列
• 其余所有行的x列值为NaN
• 下一行无需满足掩码条件

期望输出（注：按规则，正确输出应为行3的x为e，原示例输出可能存在笔误） #

a      b  c    x
0   100   1000  a  NaN
1  1123  11123  b  NaN
2   123   1123  c  NaN
3   100      0  d    e
4     1     55  e  NaN
5     0      0  f  NaN
6     1      1  g  NaN

原尝试代码 #

df.loc[mask.cumsum().eq(1) & mask, 'x'] = df.c.shift(-1)

修正后的解决方案 #

原代码逻辑基本正确，补充初始化步骤并兼容边界场景（无匹配、匹配行是最后一行）：

import pandas as pd

# 初始化DataFrame
df = pd.DataFrame(
    {
        'a': [100, 1123, 123, 100, 1, 0, 1],
        'b': [1000, 11123, 1123, 0, 55, 0, 1],
        'c': ['a', 'b', 'c', 'd', 'e', 'f', 'g'],
    }
)

# 生成掩码
mask = df.a > df.b

# 初始化x列为NaN
df['x'] = pd.NA

# 找到首次匹配的行并赋值
first_match_idx = mask.idxmax()
if mask[first_match_idx]:  # 确认存在匹配行
    if first_match_idx < len(df) - 1:  # 确保不是最后一行
        df.loc[first_match_idx, 'x'] = df.loc[first_match_idx + 1, 'c']

print(df)

测试用例覆盖 #

测试用例1：无掩码匹配场景 #

df_no_match = pd.DataFrame(
    {'a': [1,2,3], 'b': [4,5,6], 'c': ['x','y','z']}
)
mask_no_match = df_no_match.a > df_no_match.b
df_no_match['x'] = pd.NA
first_idx = mask_no_match.idxmax()
if mask_no_match[first_idx]:
    if first_idx < len(df_no_match)-1:
        df_no_match.loc[first_idx, 'x'] = df_no_match.loc[first_idx+1, 'c']
print(df_no_match)

输出：

a  b  c     x
0  1  4  x  <NA>
1  2  5  y  <NA>
2  3  6  z  <NA>

测试用例2：首次匹配行是最后一行 #

df_last_row = pd.DataFrame(
    {'a': [1,2,10], 'b': [4,5,8], 'c': ['x','y','z']}
)
mask_last = df_last_row.a > df_last_row.b
df_last_row['x'] = pd.NA
first_idx = mask_last.idxmax()
if mask_last[first_idx]:
    if first_idx < len(df_last_row)-1:
        df_last_row.loc[first_idx, 'x'] = df_last_row.loc[first_idx+1, 'c']
print(df_last_row)

输出：

a  b  c     x
0   1  4  x  <NA>
1   2  5  y  <NA>
2  10  8  z  <NA>

测试用例3：多次匹配掩码场景 #

df_multiple = pd.DataFrame(
    {'a': [5,3,6,7], 'b': [3,4,2,1], 'c': ['w','x','y','z']}
)
mask_multiple = df_multiple.a > df_multiple.b
df_multiple['x'] = pd.NA
first_idx = mask_multiple.idxmax()
if mask_multiple[first_idx]:
    if first_idx < len(df_multiple)-1:
        df_multiple.loc[first_idx, 'x'] = df_multiple.loc[first_idx+1, 'c']
print(df_multiple)

输出：

a  b  c     x
0  5  3  w     x
1  3  4  x  <NA>
2  6  2  y  <NA>
3  7  1  z  <NA>

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