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 应用于行来获取前三个值的索引，则可以将生成的索引与数据框的 columns 属性结合起来以获得您正在查找的结果。

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 方法只需要 20,000 行就花费了 1,500 倍以上的时间:

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