如何用列表筛选Pandas DataFrame且匹配列表中元素的出现次数

实现思路 #

先统计目标列表中各元素的出现频次，再对DataFrame的目标列分组后生成组内序号，筛选出组内序号不超过对应元素频次的行即可。

完整代码示例 #

import pandas as pd
from collections import Counter

# 原始数据
my_list = [1,1,2,3,4,4]
my_df = pd.DataFrame({
    'col_1': ['a','b','c','d','e','f','g','h'],
    'col_2': [1,1,2,3,3,4,4,4]
})

# 统计列表元素出现频次
count_map = Counter(my_list)
# 分组生成组内序号后筛选
result = my_df[my_df.groupby('col_2').cumcount() + 1 <= my_df['col_2'].map(count_map)]

print(result)

输出结果 #

col_1  col_2
0     a      1
1     b      1
2     c      2
3     d      3
5     f      4
6     g      4

逻辑说明 #

• groupby('col_2').cumcount()会给每个col_2分组内的行生成从0开始的递增序号，加1后转为从1开始计数
• map(count_map)会把col_2的每个值替换为它在目标列表中的出现次数
• 筛选逻辑会自动保留每个值对应的前N行，N为该值在列表中的出现次数，最终返回的行数和列表长度完全一致

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