Pandas Диаграмма DataFrame

Что такое эффективный способ получить диагональ квадрата DataFrame. Я ожидаю, что результат будет Series с MultiIndex с двумя уровнями, первый из которых является индексом DataFrame, вторым уровнем которого являются столбцы DataFrame.

Настройка

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
                 )

Я хочу увидеть это:

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

A  a    2
B  b    2
C  c    3

Ответ 1

Если вы не против использования numpy, вы можете использовать numpy.diag

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

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

Ответ 2

Вы можете сделать что-то вроде этого:

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 даст вам диагональные значения в виде массива np, вы можете затем построить мультииндекс, закрепив индекс и столбцы и передайте это как желаемый индекс в DataFrame ctor.

На самом деле сложное многоиндексное поколение не должно быть настолько сложным:

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

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

Но ответ johnchase является более аккуратным

Ответ 3

Вы также можете использовать iat в понимании списка, чтобы получить диагональ.

>>> 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