寻找一种基于 2 列过滤 Pandas Dataframe 的快速方法

Question

我正在努力改进下面的代码。

我试过使用 lambda(filter)) 但时间几乎相同。

即使可能，也无法为此创建矢量化。

df = pd.DataFrame({'A':[1,1,2,3,3,2,1],
                  'B':['foo','baa','foo','baa','foo','foo','foo']})

remove_list = [a for a in set(df['A'].values) if len(df[df['A']==a]['B'].unique())<2]

df[~df['A'].isin(remove_list)]

Answer 1

IIUC，尝试 groupby 和 nunique:

>>> df[df.groupby("A")["B"].transform('nunique').ge(2)]
   A    B
0  1  foo
1  1  baa
3  3  baa
4  3  foo
6  1  foo