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

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

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，但实际上它似乎比较慢，并且已经被弃用了。有什么建议可以让它跑得更快吗？提前谢谢

如果所有其他操作都失败，您可以对每行使用

.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

---

distance

每行的长度是相同的吗？谢谢，我马上就来试一试！numpy数组choose应该可以工作，但可能会从numpy#1开始，以防万一有一天我有超过32个（这将是一个不同的学区！）。是的，我已经编辑了

choose

-它比其他方法没有多大好处。