使用 .take() 索引多维数组

Question

我有一个形状（n，x，y）的多维数组。 对于这个例子可以使用这个数组

A = array([[[ 0,  1,  2],
            [ 3,  4,  5],
            [ 6,  7,  8],
            [ 9, 10, 11]],

           [[12, 13, 14],
            [15, 16, 17],
            [18, 19, 20],
            [21, 22, 23]],

           [[24, 25, 26],
            [27, 28, 29],
            [30, 31, 32],
            [33, 34, 35]]])

然后我有另一个多维数组，它具有我想在原始数组 A 上使用的索引值。它具​​有形状 (z,2) 并且这些值表示行值索引的

Row_values = array([[0,1],
                    [0,2],
                    [1,2],
                    [1,3]])

所以我想使用 row_values 中的所有索引值来应用 A 中的三个数组中的每一个，所以我最终得到一个形状为 (12,2,3) 的最终数组

Result = ([[[0,1,2],
            [3,4,5]],
           [[0,1,2],
            [6,7,8]],
           [[3,4,5],
            [6,7,8]]
           [[3,4,5],
            [9,10,11],
           [[12,13,14],
            [15,16,17]],
           [[12,13,14],
            [18,19,20]],
           [[15,16,17],
            [18,19,20]],
           [[15,16,17],
            [21,22,23]],
           [[24,25,26],
            [27,28,29]],
           [[24,25,26],
            [30,31,32]],
           [[27,28,29],
            [30,31,32]],
           [[27,28,29],
            [33,34,35]]]

我曾尝试使用 np.take() 但未能使其正常工作。 不确定是否有另一个更容易使用的 numpy 函数

Answer 1

我们可以利用NumPy的高级索引并将np.repeat和np.tile与它一起使用。

cidx = np.tile(Row_values, (A.shape[0], 1))
ridx = np.repeat(np.arange(A.shape[0]), Row_values.shape[0])

out = A[ridx[:, None], cidx]
# out.shape -> (12, 2, 3)

输出：

array([[[ 0,  1,  2],
        [ 3,  4,  5]],

       [[ 0,  1,  2],
        [ 6,  7,  8]],

       [[ 3,  4,  5],
        [ 6,  7,  8]],

       [[ 3,  4,  5],
        [ 9, 10, 11]],

       [[12, 13, 14],
        [15, 16, 17]],

       [[12, 13, 14],
        [18, 19, 20]],

       [[15, 16, 17],
        [18, 19, 20]],

       [[15, 16, 17],
        [21, 22, 23]],

       [[24, 25, 26],
        [27, 28, 29]],

       [[24, 25, 26],
        [30, 31, 32]],

       [[27, 28, 29],
        [30, 31, 32]],

       [[27, 28, 29],
        [33, 34, 35]]])