python - 从两个 pandas DataFrame 中选择前 N 个

Question

假设有两个 pandas DataFrame:df1 和 df2。 df1 是一个方形数据框，如下所示

import numpy as np 
import pandas as pd
item_names = [2,7,9,10,11,13,14,21,24]
np.random.seed(123)
nums = np.round(np.random.random(size=(9,9)),2)
df1 = pd.DataFrame(nums, index=item_names, columns=item_names)

df1 输出:

      2     7     9     10    11    13    14    21    24
2   0.70  0.29  0.23  0.55  0.72  0.42  0.98  0.68  0.48
7   0.39  0.34  0.73  0.44  0.06  0.40  0.74  0.18  0.18
9   0.53  0.53  0.63  0.85  0.72  0.61  0.72  0.32  0.36
10  0.23  0.29  0.63  0.09  0.43  0.43  0.49  0.43  0.31
11  0.43  0.89  0.94  0.50  0.62  0.12  0.32  0.41  0.87
13  0.25  0.48  0.99  0.52  0.61  0.12  0.83  0.60  0.55
14  0.34  0.30  0.42  0.68  0.88  0.51  0.67  0.59  0.62
21  0.67  0.84  0.08  0.76  0.24  0.19  0.57  0.10  0.89
24  0.63  0.72  0.02  0.59  0.56  0.16  0.15  0.70  0.32

df2存储item及其对应的group信息如

df2 = pd.DataFrame({'item': item_names,
                    'group':['a1','a1','a1','a2',
                             'a2','a2','a2','a3','a3']})

df2 输出:

   item group
0     2    a1
1     7    a1
2     9    a1
3    10    a2
4    11    a2
5    13    a2
6    14    a2
7    21    a3
8    24    a3

目标是编写一个函数，它可以使用这两个 DataFrame 根据相应的值(最大)选择特定行(项目名称)中的前 N ​​个项目' 信息。但是，返回的前 N ​​个项目和查询项目全部必须来自“不同组”。比如

查询项 (item = 10) 位于 df1 (item = 10) 的第 4 行。返回的前 2 项将是 [9, 21] 而不是 [9, 14]。因为，项目 10 来自组 = a2，并且任何返回的项目(前 N 个)不应来自 a2 组。我查过Scott Boston solution对于类似的问题，但它无法避免前 N 个项目和查询项目来自同一组。有什么建议么？非常感谢

Answer 1

IIUC，您想要选择除同一组中的值之外的 N 个最大值。

这是一个执行此操作的函数:

def get_top_N(idx, N=2):
    group = df2.set_index('item')['group']
    incl = group[group.ne(group[idx])].index
    return df1.loc[idx, incl].nlargest(2).index.to_list()
    
get_top_N(10)
# [9, 21]

如果您还想确保所有值都来自不同的组(不清楚是否有要求，因为您的示例就是这种情况)。您还可以这样做:

def get_top_N_diff(idx, N=2):
    group = df2.set_index('item')['group']
    incl = group[group.ne(group[idx])].index
    s = df1.loc[idx, incl]
    return s.sort_values(ascending=False).groupby(group).idxmax().to_list()[:N]

get_top_N(11)      # same group
# [9, 7]

get_top_N_diff(11) # different groups
# [9, 24]