首页 > 解决方案 > 如何将组数映射到熊猫中的另一个数据框

问题描述

我有 2 个数据框,如下所示:

df1 = pd.DataFrame({'A': ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H'],
                    'B': ['C1', 'C1', 'C1', 'C2', 'C2', 'C2', 'C2', 'C2'],
                    'rank': [2, 5, 1, 8, 6, 3, 4, 7]})

Out[3]: 
   A   B  rank
0  A  C1     2
1  B  C1     5
2  C  C1     1
3  D  C2     8
4  E  C2     6
5  F  C2     3
6  G  C2     4
7  H  C2     7

df2 = pd.DataFrame({'B': ['C1', 'C1', 'C1', 'C2'],
                    'C': [1, 2, 3, 4]})

Out[6]: 
    B  C
0  C1  1
1  C1  2
2  C1  3
3  C2  4

我想选择 df1 中排名最高的 3 行(按“排名”列),但每组最多只能选择 4 个名称(B 列),这需要包括 df2 中每组中的行数。

生成的数据框应如下所示:

   A   B  rank
2  C  C1     1
5  F  C2     3
6  G  C2     4

逻辑:

C1 组的 df2 中的行数为 3(在 df1 中最多可从该组中选择 1 行),C2 的计数为 1(最多可从 df1 中选择 3 行)

项目 C 排名最高,因此被选中,现在 C1 组的总数为 4 项目 F 和项目 G 是排名第二的项目,并且是 C2 组的一部分,总数为 3,因此少于 4

我尝试了以下方法:

df1.sort_values('rank').groupby('B').head(4).head(5)

但这仅限于在 B 中选择最多 4 行组,仅选择 df1 中的行并忽略 df2

标签: pythonpandaspandas-groupby

解决方案


这是一个想法:

max_per_group = 4

# maximal rows to pick from each group
max_sizes = max_per_group - df2.groupby('B').size()

# 4 rows from each group
heads = df1.sort_values('rank').groupby('B').head(max_per_group)

# enumerate the rows within each group
enum = heads.groupby('B').cumcount()

# output
heads[enum<heads['B'].map(sizes).fillna(max_per_group)].head(3)

输出:

   A   B  rank
2  C  C1     1
5  F  C2     3
6  G  C2     4

推荐阅读