Python 如何在巨大数据帧的每一行中找到top-n值的列索引

我有一个数据帧的格式：（示例数据）

其中，分数范围介于[0,1]和我希望生成一个函数，该函数针对每个id（行）计算前n个度量，其中n是函数的输入以及原始数据帧

我的理想输出是：（例如n=3）

现在，我已经编写了一个可以工作的函数：

def top_n_partners(scores,top_n=3):
metrics = np.array(scores.columns)
records=[]
for rec in scores.to_records():
    rec = list(rec)
    ID = rec[0]
    score_vals = rec[1:]
    inds = np.argsort(score_vals)
    top_metrics = metrics[inds][::-1]
    dic = {
        'top_score_%s' % (i+1):top_metrics[i]
        for i in range(top_n)
    }
    dic['ID'] = ID
    records.append(dic)
top_n_df = pd.DataFrame(records)
top_n_df.set_index('ID',inplace=True)
return top_n_df

然而，它似乎效率很低/速度很慢，尤其是对于我将要运行的数据量（具有数百万行的数据帧），我想知道是否有更聪明的方法来实现这一点？

您可以使用：

感谢您改进：

n = 3
df = pd.DataFrame(df.columns[df.values.argsort(1)[:,-n+2:1:-1]], 
                               index=df.index)

df = df.rename(columns = lambda x: 'Top_{}'.format(x + 1))
print (df)
      Top_1    Top_2    Top_3
ID                           
1   Metric4  Metric5  Metric1
2   Metric2  Metric3  Metric5
3   Metric3  Metric4  Metric1
4   Metric1  Metric3  Metric4                

---

使用熊猫造型的另一种方式：

df.set_index('ID', inplace=True)
df_out = df.rank(axis=1, ascending=False).astype(int).reset_index().melt(id_vars='ID').query('value <= 3').pivot(index='ID',columns='value')
df_out.columns = df_out.columns.droplevel().astype(str)
df_out = df_out.add_prefix('Top_')
print(df_out)

---

使用：

df.values.argsort（1）[：，-n+2:1:-1]

可以获得更高的性能。您可以执行

row.nlargest（3）.index

。只需将其应用于每一行。

print (np.argsort(-df.values, axis=1)[:,:3])
[[3 4 0]
 [1 2 4]
 [2 3 0]
 [0 2 3]]

print (df.columns[np.argsort(-df.values, axis=1)[:,:3]])

Index([['Metric4', 'Metric5', 'Metric1'], ['Metric2', 'Metric3', 'Metric5'],
       ['Metric3', 'Metric4', 'Metric1'], ['Metric1', 'Metric3', 'Metric4']],
      dtype='object')

df = pd.DataFrame(df.columns[np.argsort(-df.values, axis=1)[:,:3]], 
                               index=df.index)
df = df.rename(columns = lambda x: 'Top_{}'.format(x + 1))
print (df)
      Top_1    Top_2    Top_3
ID                           
1   Metric4  Metric5  Metric1
2   Metric2  Metric3  Metric5
3   Metric3  Metric4  Metric1
4   Metric1  Metric3  Metric4 

n = 3
df = pd.DataFrame(df.columns[df.values.argsort(1)[:,-n+2:1:-1]], 
                               index=df.index)

df = df.rename(columns = lambda x: 'Top_{}'.format(x + 1))
print (df)
      Top_1    Top_2    Top_3
ID                           
1   Metric4  Metric5  Metric1
2   Metric2  Metric3  Metric5
3   Metric3  Metric4  Metric1
4   Metric1  Metric3  Metric4                

df.set_index('ID', inplace=True)
df_out = df.rank(axis=1, ascending=False).astype(int).reset_index().melt(id_vars='ID').query('value <= 3').pivot(index='ID',columns='value')
df_out.columns = df_out.columns.droplevel().astype(str)
df_out = df_out.add_prefix('Top_')
print(df_out)

value    Top_1    Top_2    Top_3
ID                              
1      Metric4  Metric5  Metric1
2      Metric2  Metric3  Metric5
3      Metric3  Metric4  Metric1
4      Metric1  Metric3  Metric4