Python中的循环优化

Question

我有一个数据框（df），其中包含行驶距离，并且我已根据某些条件分配了标签。

distance=[0,0.0001,0.20,1.23,4.0]
df = pd.DataFrame(distance,columns=["distance"])
df['label']=0
for i in range(0, len(df['distance'])):   
      if (df['distance'].values[i])<=0.10:
          df['label'][i]=1
      elif (df['distance'].values[i])<=0.50:
          df['label'][i]=2
      elif (df['distance'].values[i])>0.50:
          df['label'][i]=3

这段代码运行良好。然而，我有超过100万条包含距离信息的记录，而这个for循环的执行时间比预期还要长。我们能否优化这段代码以减少执行时间？

- sshr

---

我认为如果最后的 elif... 变成 else:，你可以缩短微秒数。 - depperm

你的第二个elif应该是 0.10 < df['distance'].values[i])<=0.50？我建议为每个条件创建一个新的数据框列，然后合并它们，在切片和广播方面比循环更快。 - Andrew

两件事：如果您将df['label']设置为0，为什么df['label'][i] = 1不会创建错误？另外，我不知道您是否使用Python2或Python3 - 但是对于Python2，请用xrange替换range。 - kratenko

Answer 1

总的来说，除非绝对必要，否则不应该在DataFrames上进行循环。使用已经被优化的内置Pandas函数或使用矢量化方法通常能够得到更好的性能。
在这种情况下，您可以使用loc和Boolean indexing进行赋值操作：

# Initialize as 1 (eliminate need to check the first condition).
df['label'] = 1

# Case 1: Between 0.1 and 0.5
df.loc[(df['distance'] > 0.1) & (df['distance'] <= 0.5), 'label'] = 2

# Case 2: Greater than 0.5
df.loc[df['distance'] > 0.5, 'label'] = 3

另一个选择是使用 pd.cut。这种方法对于问题示例问题更加专业化，而布尔索引是一种更通用的方法。

# Get the low and high bins.
low, high = df['distance'].min()-1, df['distance'].max()+1

# Perform the cut.  Add one since the labels start at zero by default.
df['label'] = pd.cut(df['distance'], bins=[low, 0.1, 0.5, high], labels=False) + 1

在上面的代码中，您还可以使用labels = [1,2,3]，而不是将结果加1。这将使df ['labels']具有分类数据类型，而不是整数数据类型。根据您的用例，这可能很重要，也可能不重要。

任何一种方法的输出结果如下：

   distance  label
0    0.0000      1
1    0.0001      1
2    0.2000      2
3    1.2300      3
4    4.0000      3

- root

Answer 2

使用cut通过给箱子分配标签：

pd.cut(df.distance, [-np.inf, 0.1, 0.5, np.inf], labels=[1,2,3])

0    1
1    1
2    2
3    3
4    3

- Zeugma

Answer 3

关于在片段的副本上设置值的警告，但也许有人可以建议更清晰的替代方案？

仅基于花式索引以根据距离获取子阵列，然后将所需值写入其中。

df.loc[:, "label"][df.loc[:, "distance"] <= 0.1] = 1
df.loc[:, "label"][(0.1 < df.loc[:, "distance"]) & (df.loc[:, "distance"] <= 0.5)] = 2
df.loc[:, "label"][df.loc[:, "distance"] > 0.5] = 3

编辑：新的改进版本，不使用链式索引。

- Andrew