如何用Pandas筛选分组后唯一值数量超2的数据并生成CSV

需求 #

• 按colA分组，统计满足colC>0且colB在指定集合中的colB唯一值数量，仅保留计数大于2的分组
• 将符合条件的colA值存入集合
• 生成包含colA、colB、colC的CSV文件，要求：colB在指定集合中、colC>0、且对应的colA属于上述符合条件的集合

现有代码 #

key_id = {1008,1051,1003,1005,1072,1079,1092,1093,1004,2000}

(filtData.loc[(filtData['colC'] > 0) & (filtData['colB'].isin(key_id)),'colB']
    .groupby(filtData['colA'])
    .agg(UniqueNum='nunique')
    .reset_index())
# Above .loc might need additional booleans to include only uniqueNum greater than 2

输入数据 #

colA  colB       colC
0        1105904         1007   867.000000
1        1105904         1008    93.330002
2        1105904         2000     0.420000
3        1105904         1004     0.000000
4        1105904         1005     4.000000
5        1105905         1008     4.000000
6        1105905         1001     0.000000
7        1105905         1004     0.000000
8        1105905         1005     1.500000
9        1105905         1072     4.000000
10       1105905         1093     0.250000
11       1111111         1001     0.350000
12       1111111         1002     0.100000
13       1111111         1003     0.000000
14       1111111         1004     0.250000
15       1111111         1093     0.250000

期望统计输出 #

colA    UniqueNum
0   1105904 3
1   1105905 4

# colA value 1111111 should not show up because the unique number count is not greater than 2. 

解决方案 #

1. 完成统计并筛选计数大于2的结果 #

在现有代码基础上添加筛选步骤即可，两种常用实现方式：

方式一：用query方法（简洁直观）

key_id = {1008,1051,1003,1005,1072,1079,1092,1093,1004,2000}

filtered_result = (filtData.loc[(filtData['colC'] > 0) & (filtData['colB'].isin(key_id)),'colB']
    .groupby(filtData['colA'])
    .agg(UniqueNum='nunique')
    .reset_index()
    .query('UniqueNum > 2'))  # 筛选计数大于2的行

print(filtered_result)

方式二：用布尔索引

key_id = {1008,1051,1003,1005,1072,1079,1092,1093,1004,2000}

filtered_result = (filtData.loc[(filtData['colC'] > 0) & (filtData['colB'].isin(key_id)),'colB']
    .groupby(filtData['colA'])
    .agg(UniqueNum='nunique')
    .reset_index())

# 筛选UniqueNum大于2的结果
filtered_result = filtered_result.loc[filtered_result['UniqueNum'] > 2]

print(filtered_result)

运行后即可得到你期望的统计输出。

2. 将符合条件的colA存入集合 #

直接提取filtered_result中的colA列转成集合：

valid_colA = set(filtered_result['colA'])

3. 生成目标CSV文件 #

先筛选原数据中满足所有条件的行，再导出为CSV：

# 组合筛选条件：colB在key_id中、colC>0、colA属于valid_colA
final_data = filtData.loc[
    (filtData['colB'].isin(key_id)) & 
    (filtData['colC'] > 0) & 
    (filtData['colA'].isin(valid_colA))
]

# 导出为CSV，index=False表示不保存行索引
final_data.to_csv('target_data.csv', index=False)

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