如何优化一个函数，该函数包含for循环和数据帧中的2000万行

Posted 2021-05-07

我有一个如下所示的熊猫数据框df：

student_id   category_id  count
1              111        10
2              111        5
3              222        8
4              333        5
5              111        6

同样，我有2000万行。

我想为每个student_id计算评分。例如，让我们考虑一个category_id“ 111”。我们有3个此类的Student_id 1、2和5。student_id 1有10个计数，student_id 2有5个计数，student_id 5有6个计数。category_id的每个student_id的评分是通过以下公式计算的：

(count per student_id / total number of counts per category_id) * 5 

对于student_id 1-> 10/21 * 5 = 2.38

对于student_id 2-> 5/21 * 5 = 1.19

对于student_id 5-> 6/21 * 5 = 1.43

下面是我已经要计算的函数：

countPerStudentID = datasetPandas.groupby('student_id').agg(list)
countPerCategoryID = datasetPandas.groupby('category_id').agg(list)

studentIDMap = dict()
def func1(student_id):
    if student_id in studentIDMap:
        return studentIDMap[student_id]
    runningSum = 0
    countList = countPerStudentID.loc[student_id, 'count']
    for count in countList:
        runningSum += count
    studentIDMap[student_id] = runningSum
    return studentIDMap[student_id]

#Similar to the above function
categoryIDMap = dict()
def func2(category_id):
    if category_id in categoryIDMap:
        return categoryIDMap[category_id]
    runningSum = 0
    countList = countPerCategoryID.loc[category_id, 'count']
    for count in countList:
        runningSum += count
    categoryIDMap[category_id] = runningSum
    return categoryIDMap[category_id]

最后我从下面调用这两个函数：

#Calculating rating category-wise
rating = []
for index, row in df.iterrows():

    totalCountPerCategoryID = func1(row['category_id'])
    totalCountPerStudentID = func2(row['student_id'])

    rating.append((totalCountPerStudentID / totalCountPerCategoryID) * 5)

df['rating'] = rating

必填输出：

student_id   category_id  count   rating
1              111        10       2.38
2              111        5        1.19
3              222        8         5
4              333        5         5 
5              111        6        1.43

由于数据量巨大，运行该程序需要大量时间。我想知道如何优化此计算

提前感谢

答案

您不需要循环，这是groupby情况：

df['rating'] = df['count']/df.groupby('category_id')['count'].transform('sum') * 5

输出：

   student_id  category_id  count    rating
0           1          111     10  2.380952
1           2          111      5  1.190476
2           3          222      8  5.000000
3           4          333      5  5.000000
4           5          111      6  1.428571

另一答案

好上帝，请不要使用iterrows和append，甚至请不要使用together。难怪您的表现令人抓狂。对于pandas，iterrows应该是不得已的选择。

您应该可以使用向量化方法来实现：

>>> df['rating'] = df['count'].div(df.groupby('category_id')['count'].transform(sum)).mul(5)
>>> df
   student_id  category_id  count    rating
0           1          111     10  2.380952
1           2          111      5  1.190476
2           3          222      8  5.000000
3           4          333      5  5.000000
4           5          111      6  1.428571