python - 使用具有特定值计数的 GroupBy 过滤 Pandas DataFrame

Question

我想将 pandas DataFrame 过滤到特定行的组具有特定列值的最小计数的行。

例如，仅返回 df 的行/组，其中 ['c2','c3'] 组至少有 2 行的 'c1' 值为 1:

df = pd.DataFrame({'c1':[0,1,0,1,1,0], 'c2':[0,0,0,1,1,1], 'c3':[0,0,0,1,1,1]})

结果应该只返回索引为 3、4、5 的行，因为只有 [c2,c3] = [1,1] 组至少有 2 行的 'c1' 值为 1。

df.groupby(['c2','c3']).filter(lambda x: x['c1'].count() >= 2)

没有返回所需的结果。我需要计数专门应用于 1 的计数，而不仅仅是 'c1' 的任何值。

以下是可行的，但我不确定如何让它更符合 pythonic:

s = df.groupby(['c2','c3']).apply(lambda x: x[x['c1']==1].count() >= 2).all(axis=1)
df = df.reset_index().set_index(['c2','c3']).loc[s[s].index].reset_index().set_index(['index'])    

Answer 1

使用 groupby + transform 对 bool 系列求和，我们用它来屏蔽原始 DataFrame。

m = df['c1'].eq(1).groupby([df['c2'], df['c3']]).transform('sum').ge(2)

# Alterntively assign the column
#m = df.assign(to_sum = df.c1.eq(1)).groupby(['c2', 'c3']).to_sum.transform('sum').ge(2) 

df.loc[m]
#   c1  c2  c3
#3   1   1   1
#4   1   1   1
#5   0   1   1

使用过滤器，count 不是正确的逻辑。使用 ==(或 .eq())检查 'c1' 等于特定值的位置。对 bool 级数求和并检查过滤器的每组至少有 2 次这样的事件。

df.groupby(['c2','c3']).filter(lambda x: x['c1'].eq(1).sum() >= 2)
#   c1  c2  c3
#3   1   1   1
#4   1   1   1
#5   0   1   1

---

虽然对于小型 DataFrame 而言并不明显，但随着组数量的增加，使用 lambda 的 filter 速度非常慢。 transform 很快:

import numpy as np
np.random.seed(123)
df = pd.DataFrame({'c1':np.random.randint(1,100,1000), 'c2':np.random.randint(1,100,1000), 
                   'c3':np.random.choice([1,0], 1000)})

%%timeit
m = df['c1'].eq(1).groupby([df.c3, df.c3]).transform('sum').ge(2)
df.loc[m]
#5.21 ms ± 15.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit df.groupby(['c2','c3']).filter(lambda x: x['c1'].eq(1).sum() >= 2)
#124 ms ± 714 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)