Pandas 多索引中的访问对角

我有一个三级多索引的熊猫数据帧。我想提取前两个索引中的“对角线”。即，提取前两个指数相等的值，而第三个指数取所有可能的值

下面是生成这样一个数据帧的示例代码

import pandas as pd
import numpy as np

df = pd.DataFrame()

d = 0
e = 1
for A in xrange(3):
    for B in xrange(3):
        for C in xrange(2):
            df = df.append(dict(A=A, B=B, C=C, d=d, e=e), ignore_index=True)
            d += 2
            e += 2

df.set_index(['A', 'B', 'C'], inplace=True)
df.sortlevel(inplace=True)
print df

数据帧如下所示：

        d   e
A B C
0 0 0   0   1
    1   2   3
  1 0   4   5
    1   6   7
  2 0   8   9
    1  10  11
1 0 0  12  13
    1  14  15
  1 0  16  17
    1  18  19
  2 0  20  21
    1  22  23
2 0 0  24  25
    1  26  27
  1 0  28  29
    1  30  31
  2 0  32  33
    1  34  35

我想提取

A

和

B

相等的行（以一种有效的方式，因为真正的数据帧对于

A

和

B

有100-1000个值）。鉴于上述示例，我想提取以下数据帧：

        d   e
A B C
0 0 0   0   1
    1   2   3
1 1 0  16  17
    1  18  19
2 2 0  32  33
    1  34  35

我已经用下面的代码试过了

ind = np.diag_indices(3) + (slice(None),)
df_diag = df.loc[ind, slice(None)]
print df_diag

它打印：

        d   e
A B C
0 0 0   0   1
    1   2   3
  1 0   4   5
    1   6   7
  2 0   8   9
    1  10  11
1 0 0  12  13
    1  14  15
  1 0  16  17
    1  18  19
  2 0  20  21
    1  22  23
2 0 0  24  25
    1  26  27
  1 0  28  29
    1  30  31
  2 0  32  33
    1  34  35

正如您所看到的，它只选择了所有数据

我的问题是，如何有效地从上述数据帧中提取

A

和

B

具有相同值的行

编辑：答案的基准 使用以下代码生成更大的数据帧：

numAB = 100
numC = 10
num = numAB**2 * numC
A = np.repeat(np.arange(numAB), numAB * numC)
B = np.tile(np.repeat(np.arange(numAB), numC), numAB)
C = np.tile(np.arange(numC), numAB**2)
d = np.arange(num) * 2
e = d + 1
dflarge = pd.DataFrame(dict(A=A, B=B, C=C, d=d, e=e)).set_index(['A', 'B', 'C'])

我测试了@Alexander的答案：

%%timeit
res = pd.concat([dflarge.loc[pd.IndexSlice[n, n, :], :] for n in dflarge.index.levels[0]])

%%timeit
df2 = dflarge.reset_index()
res = df2[df2['A']==df2['B']]
res.set_index(['A', 'B', 'C'], inplace=True)

导致：

10 loops, best of 3: 96.7 ms per loop

100 loops, best of 3: 8.72 ms per loop

我测试了@Fabio的答案：

%%timeit
res = pd.concat([dflarge.loc[pd.IndexSlice[n, n, :], :] for n in dflarge.index.levels[0]])

%%timeit
df2 = dflarge.reset_index()
res = df2[df2['A']==df2['B']]
res.set_index(['A', 'B', 'C'], inplace=True)

导致：

10 loops, best of 3: 96.7 ms per loop

100 loops, best of 3: 8.72 ms per loop

---

后者的速度大约快10倍。

我首先用以下命令重置数据帧的索引：

df = df.reset_index()

然后执行筛选以选择所需内容：

d = df[df['A']==df['B']]

最后再次设置索引：

d.set_index(['A', 'B', 'C'], inplace=True)

这将返回：

        d   e
A B C        
0 0 0   0   1
    1   2   3
1 1 0  16  17
    1  18  19
2 2 0  32  33
    1  34  35

---

假设索引的第一级和第二级是相同的，您可以使用第一级的内容使用

indexlice

>>> pd.concat([df.loc[pd.IndexSlice[n, n, :], :] for n in df.index.levels[0]])
Out[278]: 
        d   e
A B C        
0 0 0   0   1
    1   2   3
1 1 0  16  17
    1  18  19
2 2 0  32  33
    1  34  35