Mutliple sums on the same dataframe

user_8098101 0
Viewed 51

I'm trying to perform multiple sums on the same dataframe and then concatenate the new dataframes into one final dataframe. Is there a concise way of doing this, or do I need use iteration ?

I have a dict of this form {key: [list_of_idx], ...} and need to groupby my dataframe for each key.

Sample data

import random
random.seed(1)

df_len = 5
df = {'idx':{i: i+1 for i in range(df_len)}, 'data':{i:random.randint(1,11) for i in range(df_len)}}
df = pd.DataFrame(df).set_index('idx')

# Groups with the idx to groupby
groups = {'a': [1,2,3,4,5],
          'b': [1,4],
          'c': [5]}

# I'm trying to avoid/find a faster way than this
dfs = []
for grp in groups:
    _df = df.loc[groups[grp]]
    _df['grp'] = grp
    _df = _df.groupby('grp').sum()

    dfs.append(_df)
dff = pd.concat(dfs)

Input (df)

   data  idx
0     2    1
1    10    2
2     9    3
3     3    4
4     6    5

Expected output (dff)

     data
grp      
a      30
c       6
b       5

Note : I'm stuck with python 2.7 and pandas 0.16.1

Time result

I tested the proposed methods and calculate the time of execution. I show the mean time per execution (using 1000 executions for each answer): I couln't test Quang Hoang first answer, because of my pandas version.

time         method  
0.00696 sec  my method (question)  
0.00328 sec  piRSquared (pd.concat)  
0.00024 sec  piRSquared (collections and defaultdict)  
0.00444 sec  Quang Hoang (2nd method : concat + reindex)
pandas-groupby pandas python-2.7 python
2 Answers

This should be (quite) a bit faster:

s = pd.Series(groups).explode()
df.reindex(s).groupby(s.index)['data'].sum()

Output:

a    30
b     5
c     6
Name: data, dtype: int64

Update: similar approach for earlier pandas version, although it might not be as fast

s = pd.concat([pd.DataFrame({'grp':a, 'idx':b}) for a,b in groups.items()],
             ignore_index=True).set_index('grp')
df.reindex(s.idx).groupby(s.index)['data'].sum()
user_42384080

Clever use of pd.concat

pd.concat({k: df.loc[v] for k, v in groups.items()}).sum(level=0)

   data
a    22
b     8
c     2

NOTE: This magically works for all columns.
Suppose we have more_data

import random
random.seed(1)
df_len = 5
df = {
    'idx':{i: i+1 for i in range(df_len)},
    'data':{i:random.randint(1,11) for i in range(df_len)},
    'more_data':{i:random.randint(1,11) for i in range(df_len)},
}
df = pd.DataFrame(df).set_index('idx')

Then

pd.concat({k: df.loc[v] for k, v in groups.items()}).sum(level=0)

   data  more_data
a    22         42
b     8         19
c     2          7

But I'd stick with more Python: collections.defaultdict

from collections import defaultdict

results = defaultdict(int)

for k, V in groups.items():
    for v in V:
        results[k] += df.at[v, 'data']

pd.Series(results)

a    22
b     8
c     2
dtype: int64

For this to work with multiple columns, I have to set up the defaultdict a tad differently:

from collections import defaultdict

results = defaultdict(lambda: defaultdict(int))

for k, V in groups.items():
    for v in V:
        for c in df.columns:
            results[c][k] += df.at[v, c]

pd.DataFrame(results)

   data  more_data
a    22         42
b     8         19
c     2          7

This is what it would look like without defaultdict but using method setdefault from the dict object instead.

results = {}

for k, V in groups.items():
    for v in V:
        for c in df.columns:
            results.setdefault(c, {})
            results[c].setdefault(k, 0)
            results[c][k] += df.at[v, c]

pd.DataFrame(results)

   data  more_data
a    22         42
b     8         19
c     2          7
user_23366540
Related