Question

我正在尝试标记（在ok）pandas DataFrame中小于'N'的所有组。我有一个有效的解决方案，但它很慢，有没有办法加快速度呢？

import pandas as pd

df = pd.DataFrame([
    [1, 2, 1],
    [1, 2, 2],
    [1, 2, 3],
    [2, 3, 1],
    [2, 3, 2],
    [4, 5, 1],
    [4, 5, 2],
    [4, 5, 3],
], columns=['x', 'y', 'z'])


keys = ['x', 'y']
N = 3

df['ok'] = True
c = df.groupby(keys)['ok'].count()
for vals in c[c < N].index:
    local_dict = dict(zip(keys, vals))
    query = ' & '.join(f'{key}==@{key}' for key in keys)
    idx = df.query(query, local_dict=local_dict).index
    df.loc[idx, 'ok'] = False
print(df)

Answer 1

使用groupby/transform/count而不是groupby/count使用builtin groupby/transform methods来形成与原始数据框df具有相同长度的系列：

c = df.groupby(keys)['z'].transform('count')

然后你可以形成一个与df具有相同长度的布尔掩码：

In [35]: c<N Out[35]: 0 False 1 False 2 False 3 True 4 True 5 False 6 False 7 False Name: ok, dtype: bool

现在，ok的分配变得更顺畅，没有循环，查询或子索引：

df['ok'] = c >= N

import pandas as pd df = pd.DataFrame([ [1, 2, 1], [1, 2, 2], [1, 2, 3], [2, 3, 1], [2, 3, 2], [4, 5, 1], [4, 5, 2], [4, 5, 3], ], columns=['x', 'y', 'z']) keys = ['x', 'y'] N = 3 c = df.groupby(keys)['z'].transform('count') df['ok'] = c >= N print(df)

产量

x y z ok 0 1 2 1 True 1 1 2 2 True 2 1 2 3 True 3 2 3 1 False 4 2 3 2 False 5 4 5 1 True 6 4 5 2 True 7 4 5 3 True

由于one of the approaches（例如transform('count')）是 Cythonized他们通常比调用groupby/transform更快使用自定义lambda函数。因此，使用
分两步计算ok列
c = df.groupby(keys)['z'].transform('count') df['ok'] = c >= N

比
快
df.assign(ok=df.groupby(keys)['z'].transform(lambda x: x.size >= N))

此外，对整个列（例如c >= N）的矢量化操作也是如此比子组上的多个操作更快。 transform(lambda x: x.size >= N))为每个群组执行一次比较x.size >= N。如果有许多小组，然后计算c >= N会提高性能。

例如，使用这个1000行的DataFrame：

import numpy as np import pandas as pd np.random.seed(2017) df = pd.DataFrame(np.random.randint(10, size=(1000, 3)), columns=['x', 'y', 'z']) keys = ['x', 'y'] N = 3

使用transform('count')的速度提高了约12倍：

In [37]: %%timeit ....: c = df.groupby(keys)['z'].transform('count') ....: df['ok'] = c >= N 1000 loops, best of 3: 1.69 ms per loop In [38]: %timeit df.assign(ok=df.groupby(keys)['z'].transform(lambda x: x.size >= N)) 1 loop, best of 3: 20.2 ms per loop In [39]: 20.2/1.69 Out[39]: 11.95266272189349

在上面的示例中有100个组：

In [47]: df.groupby(keys).ngroups Out[47]: 100

使用transform('count')的速度优势随着数量的增加而增加团体增加。例如，有955组：

In [48]: np.random.seed(2017); df = pd.DataFrame(np.random.randint(100, size=(1000, 3)), columns=['x', 'y', 'z']) In [51]: df.groupby(keys).ngroups Out[51]: 955

transform('count')方法的执行速度提高了约92倍：

In [49]: %%timeit ....: c = df.groupby(keys)['z'].transform('count') ....: df['ok'] = c >= N 1000 loops, best of 3: 1.88 ms per loop In [50]: %timeit df.assign(ok=df.groupby(keys)['z'].transform(lambda x: x.size >= N)) 10 loops, best of 3: 173 ms per loop In [52]: 173/1.88 Out[52]: 92.02127659574468

Answer 2

输入变量：

keys = ['x','y']
N = 3

使用groupby，transform和size计算是否可以：

df.assign(ok=df.groupby(keys)['z'].transform(lambda x: x.size >= N))

输出：

   x  y  z     ok
0  1  2  1   True
1  1  2  2   True
2  1  2  3   True
3  2  3  1  False
4  2  3  2  False
5  4  5  1   True
6  4  5  2   True
7  4  5  3   True

将DataFrame中的所有组标记为小于N.

2 个答案: