CODEWITHSUNDEEP
pandaspandas-groupby
user3591836
user3591836

Reputation: 993

Group dataframe by multiple columns and append the result to the dataframe

This is similar to Attach a calculated column to an existing dataframe, however, that solution doesn't work when grouping by more than one column in pandas v0.14.

For example: 

$ df = pd.DataFrame([
    [1, 1, 1],
    [1, 2, 1],
    [1, 2, 2],
    [1, 3, 1],
    [2, 1, 1]],
    columns=['id', 'country', 'source'])

The following calculation works:

$ df.groupby(['id','country'])['source'].apply(lambda x: x.unique().tolist())


0       [1]
1    [1, 2]
2    [1, 2]
3       [1]
4       [1]
Name: source, dtype: object

But assigning the output to a new column result in an error:

df['source_list'] = df.groupby(['id','country'])['source'].apply(
                               lambda x: x.unique().tolist())

TypeError: incompatible index of inserted column with frame index

Upvotes: 2

Views: 10010

Answers (3)

roman
roman

Reputation: 117345

Merge grouped result with the initial DataFrame:

>>> df1 = df.groupby(['id','country'])['source'].apply(
             lambda x: x.tolist()).reset_index()

>>> df1
  id  country      source
0  1        1       [1.0]
1  1        2  [1.0, 2.0]
2  1        3       [1.0]
3  2        1       [1.0]

>>> df2 = df[['id', 'country']]
>>> df2
  id  country
1  1        1
2  1        2
3  1        2
4  1        3
5  2        1

>>> pd.merge(df1, df2, on=['id', 'country'])
  id  country      source
0  1        1       [1.0]
1  1        2  [1.0, 2.0]
2  1        2  [1.0, 2.0]
3  1        3       [1.0]
4  2        1       [1.0]

Upvotes: 11

saladi
saladi

Reputation: 3253

An alternative method that avoids the post-facto merge is providing the index in the function applied to each group, e.g.

def calculate_on_group(x):
    fill_val = x.unique().tolist()
    return pd.Series([fill_val] * x.size, index=x.index)

df['source_list'] = df.groupby(['id','country'])['source'].apply(calculate_on_group)

Upvotes: 1

David O'Neill
David O'Neill

Reputation: 360

This can be achieved without the merge by reassigning the result of the groupby.apply to the original dataframe.

df = df.groupby(['id', 'country']).apply(lambda group: _add_sourcelist_col(group))

with your _add_sourcelist_col function being,

def _add_sourcelist_col(group):
    group['source_list'] = list(set(group.tolist()))
    return group

Note that additional columns can also be added in your defined function. Just simply add them to each group dataframe, and be sure to return the group at the end of your function declaration.

Edit: I'll leave the info above as it might still be useful, but I misinterpreted part of the original quesiton. What the OP was trying to accomplish can be done using,

df = df.groupby(['id', 'country']).apply(lambda x: addsource(x))

def addsource(x):
    x['source_list'] = list(set(x.source.tolist()))
    return x

Upvotes: 0

Related Questions