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

Question

Say I have a dataframe

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

And I want to create a new column with incremental values where category == 'A' . So it should be something like.

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

Currently I am able to do this with

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

Is there a better/pythonic way to do this (ie I don't have to initialize the value column with nan? And currently, this method assigns me a float type instead of integer which is what I want.

Answer 1

Column value is not necessary inicialize if default values is NaN , if create without [] , also for count values of mask is used 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

If need also set to nullable integers:

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

Another way could be

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