如何 index/slice 3D numpy 数组

How to index/slice 3D numpy array

我对 python/numpy 比较陌生。我有一个 TxNxN 的 3D numpy 数组。它包含一系列对称的 NxN 矩阵。我想将它转换为 TxM 的二维数组(其中 M = N(N+1)/2)。我怎样才能做到这一点?我当然可以使用 3 个循环,但我认为在 python/numpy.

中可能有更好的方法来做到这一点

看来你想得到每个对称矩阵的上三角或下三角。一个简单的方法是生成一个掩码数组并将其应用于每个二维数组:

>>> e
array([[[0, 1, 2, 3],
        [1, 2, 3, 0],
        [2, 3, 0, 1],
        [3, 0, 1, 2]],

       [[1, 2, 3, 4],
        [2, 3, 4, 1],
        [3, 4, 1, 2],
        [4, 1, 2, 3]],

       [[2, 3, 4, 5],
        [3, 4, 5, 2],
        [4, 5, 2, 3],
        [5, 2, 3, 4]]])
>>> ii, jj = np.indices(e.shape[1:])
>>> jj >= ii
array([[ True,  True,  True,  True],
       [False,  True,  True,  True],
       [False, False,  True,  True],
       [False, False, False,  True]])
>>> e[:, jj >= ii]
array([[0, 1, 2, 3, 2, 3, 0, 0, 1, 2],
       [1, 2, 3, 4, 3, 4, 1, 1, 2, 3],
       [2, 3, 4, 5, 4, 5, 2, 2, 3, 4]])

使用numpy.triu_indices函数可以做得更好,但是不能将得到的索引元组直接放在方括号之间。您需要先解压它们:

>>> i, j = np.triu_indices(e.shape[1])
>>> e[:, i, j]
array([[0, 1, 2, 3, 2, 3, 0, 0, 1, 2],
       [1, 2, 3, 4, 3, 4, 1, 1, 2, 3],
       [2, 3, 4, 5, 4, 5, 2, 2, 3, 4]])