如何在Pandas中基于number列的非NA值生成分组标签列？

需求说明 #

需要基于DataFrame的number列创建新列group，规则如下：

• 第一个非NA值之前的所有NaN对应的group为0
• 每个非NA整数值所在行，以及后续直到下一个非NA值出现前的所有NaN，对应的group为该非NA整数值

示例数据代码 #

import pandas as pd
import numpy as np

number = [np.nan, np.nan, 1, np.nan, np.nan, np.nan, 2, np.nan, np.nan, 3, np.nan, np.nan, np.nan, np.nan, np.nan, 4, np.nan, np.nan]
df = pd.DataFrame(columns=['number'])
df = df.assign(number=number)

期望输出 #

number  group
0      NaN      0
1      NaN      0
2      1.0      1
3      NaN      1
4      NaN      1
5      NaN      1
6      2.0      2
7      NaN      2
8      NaN      2
9      3.0      3
10     NaN      3
11     NaN      3
12     NaN      3
13     NaN      3
14     NaN      3
15     4.0      4
16     NaN      4
17     NaN      4

解决方案 #

利用pandas的向量化操作可高效实现需求，无需循环，适合处理数千行的大数据量：

# 生成group列：先向前填充非NA值，再将开头的NaN填充为0，最后转为整数类型
df['group'] = df['number'].ffill().fillna(0).astype(int)

原理说明 #

1. ffill()：将每一个NaN填充为其最近的前一个非NA值，让每个非NA值后面的NaN都继承该值的标记
2. fillna(0)：将第一个非NA值之前未被填充的NaN统一标记为0
3. astype(int)：将原本的浮点型结果转为整数类型，匹配分组标记的预期格式

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