python之用 Pandas 按组对唯一值求和

我有一个这样的数据框:

data = { 
    'YEAR' :    [2018,2018,2017,2018,2018,2018], 
    'SEASON':   ['SPRING', 'SPRING', 'WINTER', 'SPRING', 'SPRING', 'SPRING'], 
    'CODE':     ['A', 'A', 'A', 'B', 'C', 'D'], 
    'BUDGET':   [500,200,300,4000,700,0], 
    'QUANTITY': [1000,1000,1000,2000,300,4000] 
} 
 
df = pd.DataFrame(data) 
 
''' 
   BUDGET CODE  QUANTITY  SEASON  YEAR 
0     500    A      1000  SPRING  2018 
1     200    A      1000  SPRING  2018 
2     300    A      1000  WINTER  2017 
3    4000    B      2000  SPRING  2018 
4     700    C       300  SPRING  2018 
5       0    D      4000  SPRING  2018 
''' 

对于每个CODE，我都得到了正确的BUDGET 数量，但是不幸的是，在QUANTITY 列中我得到了每个 [Year, Season] 内该代码的总量。

我正在开发一个函数，用于在不同级别聚合我的数据框作为输入:例如，我给函数一个列表，如

my_list = [ 
  ['YEAR'] 
  ['YEAR', 'SEASON'] 
] 

并且该函数将输出一系列按每个子列表分组的数据帧。

这里的问题是我可以将 CODE 与 pd.Series.nunique 相加，并且可以求和 BUDGET 列，但如果我求和 QUANTITY 列，我显然会得到比我想要的更多的总和。我需要的是 YEAR、SEASON、CODE 的某种 sumUniques 函数。

def sumUniques(x): 
    return '???' 
 
print(df.groupby(['YEAR', 'SEASON']).agg({ 
    'CODE': pd.Series.nunique, 
    'BUDGET': sum, 
    'QUANTITY' : sumUniques 
})) 
 
''' 
             CODE  BUDGET QUANTITY 
YEAR SEASON                        
2017 WINTER     1     300      ??? 
2018 SPRING     4    5400      ??? 
 
 
--> EXPECTED RESULT: 
             CODE  BUDGET QUANTITY 
YEAR SEASON                        
2017 WINTER     1     300      1000 
2018 SPRING     4    5400      7300 
 
''' 

我在问自己哪一个可能是实现这一目标的最佳方式，我想到了 Zero's answer to "Pandas: sum values from column to unique values" :我已经试过了，似乎要么我没有正确应用它，要么没有解决我的问题，因为它引发了一个关键错误:

print(df.groupby(['YEAR', 'SEASON']).agg({ 
    'CODE': pd.Series.nunique, 
    'BUDGET': sum, 
    'QUANTITY' : lambda x: x.groupby('CODE').QUANTITY.first().sum() 
})) 
 
''' 
KeyError: 'CODE' 
''' 

我想知道完成这项工作的最佳方法是什么，希望这对其他人也有帮助!

请您参考如下方法:

根据您的意见，需要稍微复杂一些的程序才能获得您的结果。 QUANTITY 的解决方案与 jezrael 对 apply 的回答非常相似，所以感谢他。

df 
 
   BUDGET CODE  QUANTITY  SEASON  YEAR 
0     500    A      1000  SPRING  2018 
1     200    A      1000  SPRING  2018 
2     300    A      1000  WINTER  2017 
3    4000    B      2000  SPRING  2018 
4     700    C       300  SPRING  2018 
5       0    D      4000  SPRING  2018 
6     500    E      1000  SPRING  2018 
 
f = { 
        'CODE' : 'nunique',  
        'BUDGET' : 'sum' 
} 
 
g = df.groupby(['YEAR', 'SEASON']) 
v1 = g.agg(f) 
v2 = g.agg(lambda x: x.drop_duplicates('CODE', keep='first').QUANTITY.sum()) 
 
df = pd.concat([v1, v2.to_frame('QUANTITY')], 1) 
 
df 
 
             CODE  BUDGET  QUANTITY 
YEAR SEASON                         
2017 WINTER     1     300      1000 
2018 SPRING     5    5900      8300