Sum unique values by group with pandas

Question

I got a dataframe like this:

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
'''

For each CODE I got its correct BUDGET quantity, but unfortunately on the QUANTITY column I got the total quantity for that code within every [Year, Season].

I'm working on a function to aggregate my dataframe at different levels given as input: for example I'm giving the function a list like

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

and the function will output a series of dataframes grouped by each sublist.

The problem here is I can aggregate CODE with pd.Series.nunique and I can sum the BUDGET column, but if I sum also the QUANTITY column I will obviously sum up more than I want. What I would need is some sort of sumUniques function by uniques YEAR, SEASON, CODE.

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

'''

I'm asking myself which could be the best way to achieve this and I came upon Zero's answer to "Pandas: sum values from column to unique values": I've tried it out, by it seems like either I'm not applying it correctly or that does not apply my problem, as it raises a key error:

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

'''
KeyError: 'CODE'
'''

I wonder what's the best way to make this work, hope this will be of help for others too!

user_4909087 · Answer

Based on your comments, a slightly more involved procedure is required to get your result. The solution for QUANTITY is very similar to how it is in jezrael's answer with apply, so thanks to him.

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

user_2901002 · Answer

Use groupby + apply with a custom function:

def f(x):
   a = x['CODE'].nunique()
   b =  x['BUDGET'].sum()
   c = x.drop_duplicates('CODE').QUANTITY.sum()
   #Or:
   #c = x.groupby('CODE').QUANTITY.first().sum()
   return pd.Series([a,b,c], index=['CODE','BUDGET','QUANTITY'])


print (df.groupby(['YEAR', 'SEASON']).apply(f) )

             CODE  BUDGET  QUANTITY
YEAR SEASON                        
2017 WINTER     1     300      1000
2018 SPRING     4    5400      7300

Another solution:

df1 = df.groupby(['YEAR', 'SEASON']).agg({ 'CODE' : 'nunique', 'BUDGET' : 'sum'})
s = df.drop_duplicates(['YEAR', 'SEASON','CODE']).groupby(['YEAR', 'SEASON'])['QUANTITY'].sum()

df = df1.join(s.rename('QUANTITY'))
print (df)
             BUDGET  CODE  QUANTITY
YEAR SEASON                        
2017 WINTER     300     1      1000
2018 SPRING    5900     5      8300

Timings:

np.random.seed(123)
N = 1000000
a = ['WINTER', 'AUTUMN', 'SUMMER', 'SPRING']
b = list('ABCDEFGHIJKL')
c = range(1990, 2018)

data = {
    'YEAR' :    np.random.choice(c, N),
    'SEASON':   np.random.choice(a, N),
    'CODE':     np.random.choice(b, N),
    'BUDGET':    np.random.randint(1000,size= N),
    'QUANTITY': np.random.randint(1000,size= N)
}

df = pd.DataFrame(data)
print (df.head())
   BUDGET CODE  QUANTITY  SEASON  YEAR
0      92    L        95  SUMMER  2003
1     961    A       696  SPRING  1992
2     481    G       351  WINTER  1992
3     296    A        51  SPRING  1996
4     896    G        58  AUTUMN  2007

def cols(df):
    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())

    return pd.concat([v1, v2.to_frame('QUANTITY')], 1)

def jez2(df):
    df1 = df.groupby(['YEAR', 'SEASON']).agg({ 'CODE' : 'nunique', 'BUDGET' : 'sum'})
    s = df.drop_duplicates(['YEAR', 'SEASON','CODE']).groupby(['YEAR', 'SEASON'])['QUANTITY'].sum()
    return df1.join(s.rename('QUANTITY'))


def f(x):
   a = x['CODE'].nunique()
   b =  x['BUDGET'].sum()
   c = x.groupby('CODE').QUANTITY.first().sum()
   return pd.Series([a,b,c], index=['CODE','BUDGET','QUANTITY'])


print (df.groupby(['YEAR', 'SEASON']).apply(f))

print (jez2(df))
print (cols(df))

In [46]: %timeit (df.groupby(['YEAR', 'SEASON']).apply(f))
1 loop, best of 3: 674 ms per loop

In [47]: %timeit (jez2(df))
1 loop, best of 3: 1.31 s per loop

In [48]: %timeit (cols(df))
1 loop, best of 3: 1.88 s per loop

Sum unique values by group with pandas

2 Answers