如何优雅地为Pandas DataFrame数值及相邻行标记1-0标签？

问题需求 #

有一个包含数值和NaN的number列，需要添加label列，用1标记数值所在行及其上下相邻行（可扩展到自定义上下行数，比如上方1行+下方2行），其余行标记为0。预期示例：

原始解决方案（存在局限性） #

当前实现的代码虽能完成基础需求，但无法灵活扩展标记范围，代码如下：

import numpy as np
import pandas as pd
pd.set_option('display.max_rows', 100)

# 生成含NaN的DataFrame
df = pd.DataFrame(np.random.randn(100), columns=['number'])
df.loc[df.number < 1] = np.nan

# 累加当前行、上下行的非NaN值
df['label'] = df.number.fillna(0) \
            + df.number.shift(1).fillna(0) \
            + df.number.shift(-1).fillna(0)

# 将非零值转为1
df.loc[df.label > 0, 'label'] = 1
print(df)

优雅且可扩展的解决方案 #

方法1：布尔掩码扩展（推荐，向量化操作高效） #

核心思路是先标记所有非NaN行，再通过范围扩展将标记扩散到上下指定行数。这种方法支持自定义上方n行和下方m行，且完全向量化，效率远高于apply。

基础扩展版本（自定义上下行数）

import numpy as np
import pandas as pd
pd.set_option('display.max_rows', 100)

# 生成测试数据
df = pd.DataFrame(np.random.randn(100), columns=['number'])
df.loc[df.number < 1] = np.nan

# 1. 标记原始非NaN行
has_value = ~df['number'].isna()

# 2. 自定义需要扩展的上下行数
upper_rows = 1  # 上方1行
lower_rows = 1  # 下方1行

# 3. 扩展标记范围
label_mask = has_value.copy()
# 向上扩展到指定行数
for i in range(1, upper_rows + 1):
    label_mask |= has_value.shift(i, fill_value=False)
# 向下扩展到指定行数
for i in range(1, lower_rows + 1):
    label_mask |= has_value.shift(-i, fill_value=False)

# 4. 将布尔掩码转为0/1格式
df['label'] = label_mask.astype(int)
print(df)

简洁滚动窗口版本（适用于上下对称范围）

如果需要标记的上下行数对称（比如上下各1行），可以用rolling窗口的max方法实现更简洁的代码：

# 定义对称窗口大小（上下各1行+当前行，共3行）
window_size = 3
# 滚动窗口取最大值，将范围内的行标记为True
label_mask = has_value.rolling(window_size, center=True, min_periods=1).max().astype(bool)
df['label'] = label_mask.astype(int)

方法2：使用apply（不推荐，效率低） #

如果一定要用apply，可以逐行判断当前行是否在非NaN行的指定范围内，但这种方法是逐行循环，数据量大时效率极低，仅作演示：

import numpy as np
import pandas as pd
pd.set_option('display.max_rows', 100)

df = pd.DataFrame(np.random.randn(100), columns=['number'])
df.loc[df.number < 1] = np.nan

# 获取所有非NaN行的索引
value_indices = df[~df['number'].isna()].index.tolist()
upper_rows = 1
lower_rows = 1

# 用apply逐行判断
df['label'] = df.apply(
    lambda row: 1 if any(
        idx - upper_rows <= row.name <= idx + lower_rows 
        for idx in value_indices
    ) else 0,
    axis=1
)
print(df)

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