Pandas DataFrame对角线

Question

如何高效地获取一个正方形 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

- piRSquared

Answer 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

- johnchase

Answer 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来构造multiindex。

实际上，复杂的multiindex生成并不需要那么复杂：

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的答案更加简洁。

- EdChum

---

你比我先说出来了，numpy.diag 是一个好的解决方案。但是必须传递一个多级索引吗？这个方法可行吗？pd.Series(np.diag(df), index=[df.index, df.columns]) - johnchase

@johnchase 实际上这是更好的解决方案，我只是决定按照 OP 要求的字面意思去做。 - EdChum

@johnchase请提交一个带有答案的回复。我刚刚运行了它，它可以工作。 - piRSquared

Answer 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

- Alexander