python - 应用条件分组

Question

我有一个数据框，其中包含两个月值列，分别为“month1”和“month2”。如果“month1”列中的值不是“NA”，则根据“month1”列对相应的“amount”值求和。如果“month1”列中的值为“NA”，则将“month2”列中相应的“amount”值相加。

import pandas as pd
df = pd.DataFrame({'month1': [1,2,'NA', 1, 4, 'NA', 'NA'],
          'month2': ['NA',5,1, 2, 'NA', 1, 3],
          'amount': [10,20,40, 50, 60, 70, 100]})

输入输出数据框如下:

输入数据框

   month1  month2  amount
0     1.0     NaN      10
1     2.0     5.0      20
2     NaN     1.0      40
3     1.0     2.0      50
4     4.0     NaN      60
5     NaN     1.0      70
6     NaN     3.0     100

输出数据框

Answer 1

因为您的 NA 值是字符串，您可以简单地对两列进行 groupby:

# ignore month2 if month1 is NA
df.loc[df.month1.ne('NA'), 'month2'] = 'NA'

# groupby and sum
df.groupby(['month1','month2']).amount.transform('sum')

如果你不想改变你的数据，你可以这样做

s = np.where(df.month1.ne('NA'), 'NA', df['month2'])

df.groupby(['month1', s]).amount.transform('sum')

输出:

0     60
1     20
2    110
3     60
4     60
5    110
6    100
Name: amount, dtype: int64