python - 从列号获取 Pandas 列名

Question

我有 pandas 数据框，我想返回具有三个最高值的列的名称。例如:

import numpy as np
import pandas as pd

a = np.array([[2., 1., 0., 5., 4.], [6., 10., 7., 1., 3.]])
df = pd.DataFrame(a, columns=['A', 'B', 'C', 'D', 'E'])

给予:

   A   B  C  D  E
0  2   1  0  5  4
1  6  10  7  1  3

对于每一行，我想添加三个新列，其中列名的三个值最高:

   A   B  C  D  E First Second Third
0  2   1  0  5  4     D      E     A
1  6  10  7  1  3     B      C     A

我已经使用 argpartition 获取每行中前三列的索引:

inx = df.apply(np.argpartition, args=(-3,), axis=1).ix[:, -3:].values

然后需要排序

sorted_inx = inx.sort()

我不清楚我将如何获取这些列索引、获取名称，然后将它们作为三个新列填充回 df

Answer 1

虽然 Ed 的回答非常好并且在某些情况下 apply 可能是必不可少的，但我尽量避免在 pandas 中使用 apply 并完全使用矩阵运算，因为它通常会产生更好的性能。

在这种情况下，如果您使用应用于行的 numpy argsort 获得前三个值的索引，则生成的索引可以与数据框的列属性组合以获得您正在寻找的结果。

pd.concat((df, pd.DataFrame(df.columns[np.argsort(df.values, axis=1)[:, -3:][:, ::-1]], 
          columns=['First', 'Second', 'Third'])), axis=1)

   A   B  C  D  E First Second Third
0  2   1  0  5  4     D      E     A
1  6  10  7  1  3     B      C     A

虽然由于开销，给定示例的性能改进很小:

>>> %timeit pd.concat((df, pd.DataFrame(df.columns[np.argsort(df.values, axis=1)[:, -3:][:, ::-1]], columns=['First', 'Second', 'Third'])), axis=1)
100 loops, best of 3: 1.33 ms per loop

>>> %timeit df.apply(lambda x: pd.Series(x.sort_values(ascending=False).index[:3]), axis=1)
100 loops, best of 3: 3.55 ms per loop

当您将问题扩大时，改进变得非常显着，apply 方法只需要 1,500 倍的时间来处理仅 20,000 行:

a = np.array([[2., 1., 0., 5., 4.], [6., 10., 7., 1., 3.]] * 10000)
df = pd.DataFrame(a, columns=['A', 'B', 'C', 'D', 'E'])

>>> %timeit pd.concat((df, pd.DataFrame(df.columns[np.argsort(df.values, axis=1)[:, -3:][:, ::-1]], columns=['First', 'Second', 'Third'])), axis=1)
100 loops, best of 3: 4.14 ms per loop

>>> %timeit df.apply(lambda x: pd.Series(x.sort_values(ascending=False).index[:3]), axis=1)
1 loops, best of 3: 9.47 s per loop