Animeartist
Animeartist

Reputation: 1187

Pandas: Apply custom function to groups and store result in new columns in each group

I am trying to apply a custom function to each group in a groupby object and store the result into new columns in each group itself. The function returns 2 values and I want to store these values separately into 2 columns in each group.

I have tried this:

# Returns True if all values in Column1 is different.
def is_unique(x):
    status = True
    if len(x) > 1:
        a = x.to_numpy() 
        if (a[0] == a).all():
            status = False
    return status

# Finds difference of the column values and returns the value with a message.
def func(x):
    d  = (x['Column3'].diff()).dropna()).iloc[0]
    return d, "Calculated!"

# is_unique() is another custom function used to filter unique groups.
df[['Difference', 'Message']] = df.filter(lambda x: is_unique(x['Column1'])).groupby(['Column2']).apply(lambda s: func(s))

But I am getting the error: 'DataFrameGroupBy' object does not support item assignment

I don't want to reset the index and want to view the result using the get_group function. The final dataframe should look like:

df.get_group('XYZ')


   -----------------------------------------------------------------
   |   Column1 | Column2 | Column3  |  Difference   |    Message   |
   -----------------------------------------------------------------
   | 0   A     |   XYZ   |   100    |               |              |
   ----------------------------------               |              |
   | 1   B     |   XYZ   |    20    |      70       |  Calculated! |
   ----------------------------------               |              |
   | 2   C     |   XYZ   |    10    |               |              |
   -----------------------------------------------------------------

What is the most efficient way to achieve this result?

Upvotes: 4

Views: 2613

Answers (1)

jezrael
jezrael

Reputation: 862611

I think you need:

def func(x):
    d  = (x['Column3'].diff()).dropna()).iloc[0]
    last = x.index[-1]
    x.loc[last, 'Difference'] = d
    x.loc[last, 'Message'] = "Calculated!"
    return x

df1 = df.filter(lambda x: is_unique(x['Column1']))

df1 = df1.groupby(['Column2']).apply(func)

Upvotes: 3

Related Questions