如何通过Boolean Indexing为Pandas DataFrame生成对应新列

首先构造原始DataFrame：

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'col1': ['X1', 'foo', 'foo', 'X2', 'foo', 'foo', 'X3', 'foo', 'foo'],
    'col2': [np.nan, 'bar', 'bar', np.nan, 'bar', 'bar', np.nan, 'bar', 'bar'],
    'col3': [np.nan, 'baz', 'baz', np.nan, 'baz', 'baz', np.nan, 'baz', 'baz']
})

实现方案 #

我们需要将X1/X2/X3这类分组标记行的col1值，填充到后续数据行的新列new中，同时移除标记行，以下是两种简洁实现方式：

方法1：布尔标记+向前填充

# 提取分组标记行的col1值，其他行设为NaN
df['new'] = df.where(df.notna().sum(axis=1) == 1, np.nan)['col1']

# 向前填充，让后续数据行继承最近的分组标记
df['new'] = df['new'].ffill()

# 过滤掉分组标记行（保留col2非空的数据行）
result = df[df['col2'].notna()].reset_index(drop=True)

执行后得到目标结果：

col1 col2 col3 new
0  foo  bar  baz  X1
1  foo  bar  baz  X1
2  foo  bar  baz  X2
3  foo  bar  baz  X2
4  foo  bar  baz  X3
5  foo  bar  baz  X3

方法2：分组ID映射

如果需要更直观的分组逻辑，可以通过生成分组ID来实现：

# 累计分组标记行的数量，生成唯一分组ID
df['group_id'] = (df.notna().sum(axis=1) == 1).cumsum()

# 提取每个分组对应的标记值
group_labels = df[df.notna().sum(axis=1) == 1].set_index('group_id')['col1']

# 通过分组ID映射标记值到新列
df['new'] = df['group_id'].map(group_labels)

# 过滤数据行并清理临时列
result = df[df['col2'].notna()].drop('group_id', axis=1).reset_index(drop=True)

两种方法都能达成需求，第一种代码更简洁，第二种逻辑更易理解，可根据实际场景选择。

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