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

Question

熊猫apply/map是我的克星,甚至在小数据集上也可能令人痛苦地慢.下面是一个非常简单的例子,其速度差异近3个数量级.下面我创建一个Series包含100万个值的数据,只想将大于.5的值映射到"是",将小于.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 ms

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')