如何在Pandas中按ID分组后筛选合并列唯一值数量<3的组

原始数据 #

我们有如下DataFrame：

ID  from   to
A    0x    0c
A    0x    0f
A    0f    0n
B    0f    0c
B    0c    0f
C    0k    0j
C    0j    0k
C    0k    0a

需求 #

按ID分组后，仅保留**from和to两列合并后唯一值总数小于3**的分组，最终保留这些分组的所有行。

解决代码 #

可以通过groupby的filter方法直接对分组进行过滤，逻辑更简洁：

import pandas as pd

# 构造原始DataFrame
df = pd.DataFrame({
    'ID': ['A', 'A', 'A', 'B', 'B', 'C', 'C', 'C'],
    'from': ['0x', '0x', '0f', '0f', '0c', '0k', '0j', '0k'],
    'to': ['0c', '0f', '0n', '0c', '0f', '0j', '0k', '0a']
})

# 筛选符合条件的分组
result_df = df.groupby('ID').filter(
    lambda group: pd.concat([group['from'], group['to']]).nunique() < 3
)

# 输出结果
print(result_df)

运行结果 #

执行后得到目标DataFrame：

ID from  to
3  B   0f  0c
4  B   0c  0f
5  C   0k  0j
6  C   0j  0k
7  C   0k  0a

内容的提问来源于stack exchange，提问作者Olive