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的答案非常有效并且在某些情况下应用是必不可少的，但我尽量避免在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方法只需要超过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

从列号中获取Pandas列名

1 个答案: