从每一行中选择特定的列

debugcn 发表于 Dev

普里坦

我有A一个形状矩阵(p, q, r, r)和另一个I形状矩阵(r, s)。我只想从的最后一个维度中选择r个元素中的s个，以A使新矩阵的形状变为(p, q, r, s)。

为了简化（忽略前两个维度），让

>>> A
array([[5, 2, 5, 7],
       [2, 6, 4, 3],
       [4, 2, 3, 9],
       [6, 2, 4, 3]])
>>> I
array([[1, 2],
       [2, 2],
       [3, 1],
       [2, 1]])

我想要矩阵

array([[2, 5],
       [4, 4],
       [9, 2],
       [4, 2]])

我该怎么做？A[..., I]给出一个(4, 4, 2)矩阵，I从每一行中选择位于的元素。我可以通过解决问题

>>> c = np.arange(4)
>>> A[..., I][c, c, :]
array([[2, 5],
       [4, 4],
       [9, 2],
       [4, 2]])

但是我认为这需要大量的计算。有没有更有效的方法来解决此问题？

编辑：对于高维示例，请考虑I与之前相同，并且

A
array([[[[12, 15,  6, 12],
         [16, 16,  4, 17],
         [ 6, 19, 10,  9],
         [ 5, 11, 18, 17]],

        [[13, 12,  5,  6],
         [12,  7,  5,  4],
         [ 9, 19, 12,  4],
         [15,  4, 16,  7]],

        [[13,  6,  5, 17],
         [ 8,  4, 10,  9],
         [ 3, 13, 16,  4],
         [ 3,  3,  4,  4]]],


       [[[ 8,  3,  8, 18],
         [ 7, 11,  8,  7],
         [10,  8, 14,  9],
         [ 8, 12, 16,  5]],

        [[ 9, 10, 10,  7],
         [11,  6, 10,  6],
         [16, 19, 10, 14],
         [ 9, 13, 13, 19]],

        [[10,  8, 19, 12],
         [ 9, 10, 17, 19],
         [ 4, 11, 12, 14],
         [ 8,  5, 16, 10]]]])

预期产量：

array([[[[15,  6],
         [ 4,  4],
         [ 9, 19],
         [18, 11]],

        [[12,  5],
         [ 5,  5],
         [ 4, 19],
         [16,  4]],

        [[ 6,  5],
         [10, 10],
         [ 4, 13],
         [ 4,  3]]],


       [[[ 3,  8],
         [ 8,  8],
         [ 9,  8],
         [16, 12]],

        [[10, 10],
         [10, 10],
         [14, 19],
         [13, 13]],

        [[ 8, 19],
         [17, 17],
         [14, 11],
         [16,  5]]]]

A[...,I][..., c, c, :] 产生这个结果

亚图

由于您使用的是整数数组索引，因此您需要指定要从中选择这些列的行：

A[np.arange(A.shape[0])[:,None], I]
array([[2, 5],
       [4, 4],
       [9, 2],
       [4, 2]])

或者您也有np.take_along_axis：

np.take_along_axis(A, I, 1)

对于较大的形状数组(p, q, r, r)，请沿前轴完整切片，并以类似方式使用广播：

A[...,np.arange(A.shape[2])[:,None],I]
array([[[[15,  6],
         [ 4,  4],
         [ 9, 19],
         [18, 11]],

        [[12,  5],
         [ 5,  5],
         [ 4, 19],
         [16,  4]],
        ...

本文收集自互联网，转载请注明来源。

如有侵权，请联系[email protected] 删除。