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

时间:2014-08-26 16:56:14

标签: python pandas

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

E.g。从:

   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之类的函数,但是我对如何执行此操作感到有点迷失。

2 个答案:

答案 0 :(得分:8)

您可以在此处使用stack,这将在索引中生成包含行和列信息的系列,然后在其上调用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

答案 1 :(得分:1)

您可以使用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