CODEWITHSUNDEEP
pythonpandas
knightofni
knightofni

Reputation: 1956

Pandas Groupby : group **by** a column containing tuples

I'm trying to group by a column containing tuples. Each tuple has a different length.

I'd like to perform simple groupby operations on this column of tuples, such as sum or count.

Example :

df = pd.DataFrame(data={
                    'col1': [1,2,3,4]  ,
                    'col2': [('a', 'b'), ('a'), ('b', 'n', 'k'), ('a', 'c', 'k', 'z') ] ,
                    })

print df

outputs :

   col1          col2
0     1        (a, b)
1     2        (a, m)
2     3     (b, n, k)
3     4  (a, c, k, z)

I'd like to be able to group by col2 on col1, with for instance a sum.

Expected output would be :

   col2        sum_col1
0     a        7
1     b        4
2     c        4
3     n        3
3     m        2
3     k        7
3     z        4

I feel that pd.melt might be able to use, but i can't see exactly how.

Upvotes: 1

Views: 1636

Answers (1)

Primer
Primer

Reputation: 10302

Here is an approach using .get_dummies and .melt:

import pandas as pd
df = pd.DataFrame(data={
                    'col1': [1,2,3,4]  ,
                    'col2': [('a', 'b'), ('a'), ('b', 'n', 'k'), ('a', 'c', 'k', 'z') ] ,
                    })

value_col = 'col1'
id_col = 'col2'

Unpack tuples to DataFrame: 

df = df.join(df.col2.apply(lambda x: pd.Series(x)))

Create columns with values of tuples:

dummy_cols = df.columns.difference(df[[value_col, id_col]].columns)
dfd = pd.get_dummies(df[dummy_cols | pd.Index([value_col])])

Producing:

   col1  0_a  0_b  1_b  1_c  1_n  2_k  3_z
0     1    1    0    1    0    0    0    0
1     2    1    0    0    0    0    0    0
2     3    0    1    0    0    1    1    0
3     4    1    0    0    1    0    1    1

Then .melt it and clean variable column from prefixes:

dfd = pd.melt(dfd, value_vars=dfd.columns.difference([value_col]).tolist(), id_vars=value_col)
dfd['variable'] = dfd.variable.str.replace(r'\d_', '')
print dfd.head()

Yielding:

   col1 variable  value
0     1        a      1
1     2        a      1
2     3        a      0
3     4        a      1
4     1        b      0

And finally get your output:

dfd[dfd.value != 0].groupby('variable')[value_col].sum()

    variable
a           7
b           4
c           4
k           7
n           3
z           4
Name: col1, dtype: int64

Upvotes: 3

Related Questions