在应用中矢量化一个非常简单的 pandas lambda 函数

Question

pandas apply/map 是我的克星，即使在小型数据集上也可能非常缓慢。下面是一个非常简单的示例，其中速度差异接近 3 个数量级。下面我创建了一个具有 100 万个值的 Series，并且只想将大于 0.5 的值映射到“是”，将小于 0.5 的值映射到“否”。如何将其矢量化或显着加快速度？

ser = pd.Series(np.random.rand(1000000))

# vectorized and fast
%%timeit
ser > .5

1000 次循环，3 次中的最佳：每个循环 477 µs

%%timeit
ser.map(lambda x: 'Yes' if x > .5 else 'No')

1 个循环，3 个循环中的最佳值：每个循环 255 毫秒

Answer 1

np.where(cond, A, B) 是 A if cond else B 的矢量化等效项：

import numpy as np
import pandas as pd
ser = pd.Series(np.random.rand(1000000))
mask = ser > 0.5
result = pd.Series(np.where(mask, 'Yes', 'No'))
expected = ser.map(lambda x: 'Yes' if x > .5 else 'No')
assert result.equals(expected)

---

In [77]: %timeit mask = ser > 0.5
1000 loops, best of 3: 1.44 ms per loop

In [76]: %timeit np.where(mask, 'Yes', 'No')
100 loops, best of 3: 14.8 ms per loop

In [73]: %timeit pd.Series(np.where(mask, 'Yes', 'No'))
10 loops, best of 3: 86.5 ms per loop

In [74]: %timeit ser.map(lambda x: 'Yes' if x > .5 else 'No')
1 loop, best of 3: 223 ms per loop

---

由于这个系列只有两个值，您可以考虑使用 Categorical 代替：

In [94]: cat = pd.Categorical.from_codes(codes=mask.astype(int), categories=['Yes', 'No']); cat
Out[94]: 
[No, Yes, No, Yes, Yes, ..., Yes, No, Yes, Yes, No]
Length: 1000000
Categories (2, object): [Yes, No]

In [95]: %timeit pd.Categorical.from_codes(codes=mask.astype(int), categories=['Yes', 'No']); cat
100 loops, best of 3: 6.26 ms per loop

这不仅速度更快，而且内存效率更高，因为它避免了创建字符串数组。类别代码是映射到类别的整数数组：

In [96]: cat.codes
Out[96]: array([1, 0, 1, ..., 0, 0, 1], dtype=int8)

In [97]: cat.categories
Out[99]: Index(['Yes', 'No'], dtype='object')