2016-05-18 44 views
13

Che cosa è un modo efficace per ottenere la diagonale di un quadrato DataFrame. Mi aspetto che il risultato sia uno Series con uno MultiIndex con due livelli, il primo è l'indice dello DataFrame mentre il secondo livello sono le colonne dello DataFrame.pandas DataFrame diagonale

Setup

import pandas as pd 
import numpy as np 

np.random.seed([3, 1415]) 
df = pd.DataFrame(np.random.rand(3, 3) * 5, 
        columns = list('abc'), 
        index = list('ABC'), 
        dtype=np.int64 
       ) 

voglio vedere questo:

print df.stack().loc[[('A', 'a'), ('B', 'b'), ('C', 'c')]] 

A a 2 
B b 2 
C c 3 

risposta

14

Se non ti dispiace utilizzando NumPy è possibile utilizzare numpy.diag

pd.Series(np.diag(df), index=[df.index, df.columns]) 

A a 2 
B b 2 
C c 3 
dtype: int64 
6

Si potrebbe fare qualcosa di simile:

In [16]: 
midx = pd.MultiIndex.from_tuples(list(zip(df.index,df.columns))) 
pd.DataFrame(data=np.diag(df), index=midx) 

Out[16]: 
    0 
A a 2 
B b 2 
C c 3 

np.diag vi darà i valori diagonali come un array np , è quindi possibile costruire il multiindice zippando l'indice e le colonne e passare questo come indice desiderato il DataFrame ctor.

In realtà il complesso generazione multiindex non ha bisogno di essere così complicato:

In [18]: 
pd.DataFrame(np.diag(df), index=[df.index, df.columns]) 

Out[18]: 
    0 
A a 2 
B b 2 
C c 3 

Ma johnchase's answer è più ordinato

+1

Mi hai battuto per 'numpy.diag' è una buona soluzione. È necessario passare un mutliindex però? questo non funzionerebbe? 'pd.Series (np.diag (df), index = [df.index, df.columns])' – johnchase

+0

@johnchase in realtà è una soluzione migliore, ho appena deciso di fare letteralmente quello che l'OP chiedeva per – EdChum

+0

@johnchase, per favore invia una risposta. appena eseguito e funziona – piRSquared

3

È inoltre possibile utilizzare iat in un elenco di comprensione per ottenere la diagonale.

>>> pd.Series([df.iat[n, n] for n in range(len(df))], index=[df.index, df.columns]) 
A a 2 
B b 2 
C c 3 
dtype: int64