將DataFrame中的所有組標記為小於N

Question

我正在嘗試標記（ ok ）熊貓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/count ，還可以使用groupby/transform/count來形成一個與原始DataFrame df 相同長度的Series：

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

---

由於內置的groupby/transform方法 （例如transform('count') ）被Cython化，因此它們通常比使用自定義lambda函數調用groupby/transform更快。 因此，使用以下步驟分兩步計算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