使用pandas按列返回n个最小索引

Question

我有以下（简化）数据框：

df = pd.DataFrame({'X': [1, 2, 3, 4, 5,6,7,8,9,10],
'Y': [10,20,30,40,50,-10,-20,-30,-40,-50],
'Z': [20,18,16,14,12,10,8,6,4,2]},index=list('ABCDEFGHIJ'))

这将得到以下结果：

    X   Y   Z
A   1  10  20
B   2  20  18
C   3  30  16
D   4  40  14
E   5  50  12
F   6 -10  10
G   7 -20   8
H   8 -30   6
I   9 -40   4
J  10 -50   2

我想创建一个新的数据框，该数据框按列返回n个最小值的索引。

期望的输出（例如，3个最小值）：

   X  Y  Z
0  A  J  J
1  B  I  I
2  C  H  H

最佳方法是什么？

- Nick Morgan

Answer 1

你可以使用 apply 与 nsmallest：

n = 3
df.apply(lambda x: pd.Series(x.nsmallest(n).index))

#   X   Y   Z
#0  A   J   J
#1  B   I   I
#2  C   H   H

- Psidom

Answer 2

使用numpy.argsort可以得到更快的numpy解决方案：

N = 3
a = np.argsort(-df.values, axis=0)[-1:-1-N:-1]
print (a)
[[0 9 9]
 [1 8 8]
 [2 7 7]]

b = pd.DataFrame(df.index[a], columns=df.columns)
print (b)
   X  Y  Z
0  A  J  J
1  B  I  I
2  C  H  H

时间:

In [111]: %timeit (pd.DataFrame(df.index[np.argsort(-df.values, axis=0)[-1:-1-N:-1]], columns=df.columns))
159 µs ± 1.37 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [112]: %timeit (df.apply(lambda x: pd.Series(x.nsmallest(N).index)))
3.52 ms ± 49.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

- jezrael

---

这个解决方案现在已经不再有效。 关系矩阵计算正确，但第二步只会抛出ValueError: 不再支持多维索引（例如 obj[:, None]）。请先将其转换为 numpy 数组再进行索引。 我不知道该如何挽救这个问题... - Paul Alex

@PaulAlex - 解决方案是 b = pd.DataFrame(df.index.to_numpy()[a], columns=df.columns) - jezrael

Answer 3

首先，您想要按列对输入数据框进行排序，然后获取每列的所有索引列表，从这些索引创建一个数据框，然后从结果数据框中返回前n行。

def topN(df, n):
#first, sort dataframe per column
sort_x = df.sort_values(by = ['X'], ascending = True)
sort_y = df.sort_values(by = ['Y'], ascending = True)
sort_z = df.sort_values(by = ['Z'], ascending = True)
#now get a list of the indices of each sorted df
index_list_x = sort_x.index.values.tolist()
index_list_y = sort_y.index.values.tolist()
index_list_z = sort_z.index.values.tolist()
#create dataframe from lists
sorted_df = pd.DataFrame(
    {'sorted_x':index_list_x,
     'sorted_y':index_list_y,
     'sorted_z':index_list_z  
    })
#return the top n from the sorted dataframe
return sorted_df.iloc[0:n]

topN(df,3)

返回：

  X  Y  Z
0 A  J  J
1 B  I  I
2 C  H  H

- Bailey Freund