从pandas数组中获取N个最大值，索引和列标题保持不变

Question

假设我刚刚计算了一个相关矩阵。 使用pandas数据帧，我现在想要获得与其轴名称相关的最高相关性。

例如：

   a, b, c, d, e, f 
a, 0, 1, 2, 3, 4, 5,
b, 1, 0, 3, 4, 5, 6,
c, 2, 3, 0, 5, 6, 7,
d, 3, 4, 5, 0, 7, 8,
e, 4, 5, 6, 7, 0, 9,
f, 5, 6, 7, 8, 9, 0

得到：

e f 9
f d 8
f c 7
e d 7

等等...

我已经阅读了pandas文档并查看了groupby方法以及head之类的函数，但是我对如何执行此操作感到有点迷失。

Answer 1

你可以在这里使用stack ，它将在索引中生成一个包含行和列信息的Series，然后在其上调用nlargest ：

>>> df.stack()
a  a    0
   b    1
   c    2
   d    3
   e    4
   f    5
b  a    1
   b    0
   c    3
[etc.]
>>> df.stack().nlargest(6)
e  f    9
f  e    9
d  f    8
f  d    8
c  f    7
d  e    7
dtype: int64

Answer 2

您可以使用np.argpartition 。 在这里下降到NumPy似乎可以提高2-3倍的性能。

np.random.seed(0)
df = pd.DataFrame(np.abs(np.random.randn(500, 400)))

def jpp(df, n):
    flat_indices = np.argpartition(df.values.ravel(), -n)[-n:]
    row_idx, col_idx = np.unravel_index(flat_indices, df.values.shape)
    indices = list(zip(row_idx, col_idx))
    values = df.values[(row_idx, col_idx)]
    res_idx = pd.MultiIndex.from_tuples(indices)
    return pd.Series(values, index=res_idx).sort_values(ascending=False)

def dsm(df, n):
    return df.stack().nlargest(n)

assert jpp(df, n=1000).equals(dsm(df, n=1000))

%timeit jpp(df, n=1000)  # 4.65 ms per loop
%timeit dsm(df, n=1000)  # 12.1 ms per loop

%timeit jpp(df, n=5)     # 3.33 ms per loop
%timeit dsm(df, n=5)     # 10.1 ms per loop