numpy:选择索引相同的位置

Question

说我有一个多维矩阵

A = np.ones((5, 10, 5, 10, 2, 3))

现在，我想选择第一维的索引等于第三维的所有项目。也就是说，我想拥有

A[0, :, 0, ...]
A[1, :, 1, ...]
A[2, :, 2, ...]

等等....

这样做的一种复杂的（但我知道）的方法是在索引上创建一个网格，然后根据这些方法进行选择：

i1, i2, i3, i4, i5, i6 = np.meshgrid(range(5), range(10), range(5), range(10), range(2), range(3), indexing='ij')
A[i1 == i3].shape
Out[29]: (3000,)
A[i1 == i2].shape
Out[31]: (1500,)

但是，这是一种非常乏味的方式。是否有基于索引的更简单的选择方法？

奖励问题

我如何同时选择两个指数？即

A[(i1 == i3) and (i2 == i4)]

Answer 1

应该使用advanced-indexing（简单而高级，嗯...） -

非常简单

r = np.arange(a.shape[0])
out = a[r,:,r]

样本运行

1）3D数据：

In [380]: a
Out[380]: 
array([[[4, 6, 4, 8],
        [0, 4, 5, 2],
        [2, 5, 6, 7]],
       [[0, 0, 4, 6],
        [1, 7, 4, 3],
        [4, 2, 0, 1]]])
In [381]: a[0,:,0]
Out[381]: array([4, 0, 2])
In [382]: a[1,:,1]
Out[382]: array([0, 7, 2])
In [383]: r = np.arange(a.shape[0])
In [384]: a[r,:,r]
Out[384]: 
array([[4, 0, 2],
       [0, 7, 2]])

2）4D数据：

In [385]: a = np.random.randint(0,9,(2,2,3,4))
In [386]: a[0,:,0,:]
Out[386]: 
array([[0, 4, 1, 8],
       [1, 2, 2, 3]])
In [387]: a[1,:,1,:]
Out[387]: 
array([[7, 8, 3, 8],
       [0, 1, 4, 7]])
In [388]: r = np.arange(a.shape[0])
In [389]: a[r,:,r]
Out[389]: 
array([[[0, 4, 1, 8],
        [1, 2, 2, 3]],
       [[7, 8, 3, 8],
        [0, 1, 4, 7]]])

---

奖金部分：从两对尺寸中选择

简单地（再次）使用两个范围数组进行索引，但是正交通过将第一个轴放在另一个轴上以允许配对，然后索引，例如So -so -

r1 = np.arange(a.shape[0])[:,None]
r2 = np.arange(a.shape[1])
out = a[r1,r2,r1,r2]

样本运行 -

In [464]: a = np.random.randint(11,99,(3,2,3,2))
In [470]: for i in range(3):
     ...:     for j in range(2):
     ...:         print a[i,j,i,j]
     ...:         
34
12
30
20
93
76
In [471]: r1 = np.arange(a.shape[0])[:,None]
     ...: r2 = np.arange(a.shape[1])
     ...: out = a[r1,r2,r1,r2]
     ...: 
In [472]: out
Out[472]: 
array([[34, 12],
       [30, 20],
       [93, 76]])