pandas 应用 function 性能优化

Question

I have following df

activity	region	empPeople
12122	1101	2
23322	1233	40
22223	2323	0
...	...	...

I want to create a column RCA which takes value 1 if (empPeople/TotalEmpRegion) / (totalEmpActivity / totalEmp) > 1 and 0 otherwise. Then I will transform this df to pivot table with index=region and column=activity and values=rca.

I wrote the following function

def rca_emp(activity:str, region:str , emp:float):
    top = emp / df[df['region'] == region].empPeople.sum()
    bottom = df[df['activity'] == activity].empPeople.sum() / df.empPeople.sum()
    rca = top/bottom
    if rca > 1: 
        return 1
    else:
        return 0

Then I used apply method to create a column rca

# finding RCA
df['rca'] = df.apply(lambda x : rca_emp(activity=x['activity'] , region=x['region'] , emp=x['empPeople']) , axis=1)
# create a binary matrix
df.pivot(index='region', columns='activity', values='rca')

The issue is that apply function takes too much time (6047 seconds). I was wondering is there a faster way to accomplish this task?

Answer 1

Instead your function use GroupBy.transform with sum and create 0,1 in numpy.where :

s1 = df.groupby('activity')['empPeople'].transform('sum')
s2 = df.groupby('region')['empPeople'].transform('sum')

df['rca'] = np.where((df['empPeople'] / s2)  / (s1 / df.empPeople.sum()) > 1, 1, 0)

Testing ouput:

print (df)
   activity  region  empPeople
0     12122    1101          2
1     23322    1233         40
2     22223    2323          0
3     12122    1101          1
4     23322    1233          4
5     22223    2323          6



def rca_emp(activity:str, region:str , emp:float):
    top = emp / df[df['region'] == region].empPeople.sum()
    bottom = df[df['activity'] == activity].empPeople.sum() / df.empPeople.sum()
    rca = top /bottom
    if rca > 1: 
        return 1
    else:
        return 0


df['rca'] = df.apply(lambda x : rca_emp(activity=x['activity'] , region=x['region'] , emp=x['empPeople']) , axis=1)

---

s1 = df.groupby(['activity'])['empPeople'].transform('sum')
s2 = df.groupby(['region'])['empPeople'].transform('sum')

df['rca1'] = np.where((df['empPeople'] / s2)  / (s1 / df.empPeople.sum())  > 1, 1, 0)
print (df)
   activity  region  empPeople  rca  rca1
0     12122    1101          2    1     1
1     23322    1233         40    1     1
2     22223    2323          0    0     0
3     12122    1101          1    1     1
4     23322    1233          4    0     0
5     22223    2323          6    1     1