有没有一种方法可以对当前在数据帧中的行上迭代的代码进行矢量化？

Question

我现在有一些代码可以正常工作，但速度太慢了。我试图把数据帧中每一行的加权平方和相加。我想对操作进行矢量化——这似乎运行得更快、更快——但代码中有一个折痕挫败了我对矢量化的尝试

totalDist = 0.0
    
for index, row in pU.iterrows():
    totalDist += (row['distance'][row['schoolChoice']]**2.0*float(row['students']))

该行有“students”（一个整数）、distance（一个长度为n的numpy数组）和schoolChoice（一个小于或等于n-1的整数，用于指定我在计算距离数组中使用的元素）。基本上，我从numpy数组中提取特定于行的值。我已经使用了df.lookup，但实际上它似乎比较慢，并且已经被弃用了。有什么建议可以让它跑得更快吗？提前谢谢

Answer 1

如果所有其他操作都失败，您可以在每一行上使用.apply()

totalSum = df.apply(lambda row: row.distance[row.schoolChoice] ** 2 * row.students, axis=1).sum()

要加快速度，您可以导入numpy

totalSum = (numpy.stack(df.distance)[range(len(df.schoolChoice)), df.schoolChoice] ** 2 * df.students).sum()

numpy方法要求每行的距离长度相同，但是如果需要，可以将它们填充到相同的长度。（尽管这可能会影响任何收益。）

在150000行df上测试，如：

         distance  schoolChoice  students
0       [1, 2, 3]             0         4
1       [4, 5, 6]             2         5
2       [7, 8, 9]             2         6
3       [1, 2, 3]             0         4
4       [4, 5, 6]             2         5

时间：

         method    time
0      for loop   15.9s
1      df.apply    4.1s
2         numpy    0.7s