有没有办法对当前迭代 Pandas 数据帧中的行的代码进行矢量化？

Question

我现在有一些代码可以正常工作，但速度太慢了。我正在尝试将 Pandas 数据框中每一行的加权平方和相加。我想对这些操作进行矢量化处理——这似乎运行快得多——但代码中的一个错误使我的矢量化尝试失败了。

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的整数，指定距离数组I的哪个元素'用于计算）。基本上，我从 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 方法要求每行的距离长度相同 - 但是，如果需要，可以将它们填充到相同的长度。（尽管这可能会影响任何收益。）

在 150,000 行的 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

有没有办法对当前迭代 Pandas 数据帧中的行的代码进行矢量化？

Is there a way to vectorize code that currently iterates over rows in a Pandas dataframe?

python

iteration

optimization

pandas