用大熊猫的分组依据注释每一行吗？

Question

I have a dataframe that looks like this:

Company       Speciality      Payment
AcmeCorp      Roofing         50.00
AcmeCorp      Grounding       50.00
LolCorp       Roofing         106.00
LolCorp       Grounding       94.00

I'd like to add a percentage column like this:

Company       Speciality      Payment     Percent of Total Payment
AcmeCorp      Roofing         50.00       50
AcmeCorp      Grounding       50.00       50
LolCorp       Roofing         106.00      53
LolCorp       Grounding       94.00       47

What's the best way to do this?

I could do it messily using something like this:

df_m = df.groupby('Company').sum()
final_df = pd.merge(df, df_m, on='Company', suffixes=('Raw', 'Total))
final_df['Percent of Total Payment] = final_df['Payment Raw'] / final_df['Payment_Total']

But I wonder if there's a more efficient way.

Answer 1

Use groupby/transform to produce a column of the same length as the original DataFrame. This allows you to avoid calling pd.merge .

import numpy as np
import pandas as pd

df = pd.DataFrame({'Company': ['AcmeCorp', 'AcmeCorp', 'LolCorp', 'LolCorp'],
 'Payment': [50.0, 50.0, 106, 94.00],
 'Speciality': ['Roofing', 'Grounding', 'Roofing', 'Grounding']})

total = df.groupby('Company')['Payment'].transform('sum')
df['percent'] = df['Payment']/total
print(df)

yields

    Company  Payment Speciality  percent
0  AcmeCorp     50.0    Roofing     0.50
1  AcmeCorp     50.0  Grounding     0.50
2   LolCorp    106.0    Roofing     0.53
3   LolCorp     94.0  Grounding     0.47

---

Although

total = df.groupby('Company')['Payment'].transform('sum')
df['percent'] = df['Payment']/total

could be reduced to a one-liner,

df['percent'] = df.groupby('Company')['Payment'].transform(lambda x: x/x.sum())

because builtin operations like .transform('sum') are faster than those with custom functions (eg .transform(lambda x: x/x.sum()) ), the two-line version is faster (particularly for large DataFrames.)

And, of course, the two-line version could also be written as

df['percent'] = df['Payment'] / df.groupby('Company')['Payment'].transform('sum')

with no loss in speed, one less named variable, but perhaps a bit harder to read.

---

Here's a benchmark on a 100K-row DataFrame:

In [53]: %timeit using_transform(df)
100 loops, best of 3: 8.5 ms per loop

In [54]: %timeit using_one_liner(df)
10 loops, best of 3: 20.2 ms per loop

In [55]: %timeit orig(df)
10 loops, best of 3: 30.2 ms per loop

---

This is the setup used to perform the benchmark.

import numpy as np
import pandas as pd

N = 10**5
df = pd.DataFrame({'Company': np.random.choice(list('ABCD'), size=N),
    'Payment': np.random.randint(10, size=N),
    'Speciality': np.random.choice(list('XYZ'), size=N)})

def using_transform(df):
    total = df.groupby('Company')['Payment'].transform('sum')
    df['percent'] = df['Payment']/total
    return df

def using_one_liner(df):
    df['percent'] = df.groupby('Company')['Payment'].transform(lambda x: x/x.sum())
    return df

def orig(df):
    df_m = df.groupby('Company').sum()
    final_df = pd.merge(df, df_m, left_on='Company', right_index=True, suffixes=('_Raw', '_Total'))
    final_df['Percent of Total Payment'] = final_df['Payment_Raw'] / final_df['Payment_Total']
    return final_df