Pandas - 创建新列并根据过滤器分配值

Question

假设我有一个 dataframe

id  category
1   A        
2   A
3   B
4   C
5   A

我想创建一个包含增量值的新列，其中category == 'A' 。 所以它应该是这样的。

id  category  value
1   A         1
2   A         2
3   B         NaN
4   C         NaN
5   A         3

目前我能够做到这一点

df['value'] = pd.nan
df.loc[df.category == "A", ['value']] = range(1, len(df[df.category == "A"]) + 1)

有没有更好的/pythonic 方法来执行此操作（即我不必使用 nan 初始化值列？目前，此方法为我分配了一个 float 类型，而不是我想要的 integer。

Answer 1

如果默认值为NaN ，则列value不需要初始化，如果创建时不使用[] ，则掩码的计数值也使用sum ：

m = df.category == "A"
df.loc[m, 'value'] = range(1, m.sum() + 1)
df['value'] = df['value'].astype('Int64')

print (df)

0   1        A      1
1   2        A      2
2   3        B   <NA>
3   4        C   <NA>
4   5        A      3

如果还需要设置为可为空的整数：

m = df.category == "A"
df['value'] = m.cumsum().where(m).astype('Int64')
print (df)
   id category  value
0   1        A      1
1   2        A      2
2   3        B   <NA>
3   4        C   <NA>
4   5        A      3

Answer 2

另一种方式可能是

df['value'] = df['category'].eq('A').cumsum()
df['value'][df['category']!='A'] = pd.NA