Question

说我有一个多维矩阵

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

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

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

依旧......

一个复杂的（但我知道）方法是在索引上创建一个meshgrid，然后根据这些选择：

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

奖励部分：从两对维度中选择

简单地（再次）使用两个范围数组进行索引，但是彼此正交通过在另一个轴上获取第一个以允许配对然后索引，就像这样 -

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

Numpy：选择索引相同的位置

1 个答案: