将零重新定位到多维 numpy 数组中最后一个维度的末尾

Question

假设您有以下形状数组 (5, 4, 3)：

x = array([[[ 0,  2,  1],
            [ 3,  4,  0],
            [ 8,  0,  6],
            [ 9,  0,  0]],

           [[12,   0, -14],
            [ 0,  16,  17],
            [ 0, -19,  20],
            [21,  22,   0]],

           [[24,   0, 26],
            [27,   0, 29],
            [ 0, -31,  0],
            [33,  34,  0]],

           [[ 0, 37, 38],
            [39,  0, 41],
            [42, 43,  0],
            [45,  0, 47]],

           [[49, 48,  0],
            [51,  0, 53],
            [ 0,  0, 56],
            [ 0, 59, 58]]])

其中维度 2 的每个 "row" 中至少有 1 个零（例如 [0, 1, 2] 或 [6, 0, 8]）。是否有任何矢量化方法可以将所有零重新定位到维度 2 的末尾，同时保留非零元素的顺序。比如将上面的数组取为：

y = array([[[ 2,  1,  0],
            [ 3,  4,  0],
            [ 8,  6,  0],
            [ 9,  0,  0]],

           [[ 12, -14,  0],
            [ 16,  17,  0],
            [-19,  20,  0],
            [ 21,  22,  0]],

            ...

           [[49, 48,  0],
            [51, 53,  0],
            [56,  0,  0],
            [59, 58,  0]]])

我可以通过遍历前两个维度并重新排列维度 2 中的每个 3 元素列表来做到这一点，但实际上，我有一个更大的 3 维数组并且正在寻求加速。

编辑：这不需要就地完成。还有，在真题中，非零元素是没有排序的，是非零浮点数，包括负数。

Answer 1

一种方法使用 mask -

import numpy as np

# Mask of non-zeros elements in the output array        
mask = ~np.sort(x==0,2)

# Setup output array and put non zero input elements into it
out = np.zeros_like(x)
out[mask] = x[x!=0]

样本运行-

输入：

In [96]: x
Out[96]: 
array([[[  0,   2,   1],
        [  3,   4,   0],
        [  8,   0,   6],
        [  9,   0,   0]],

       [[ 12,   0, -14],
        [  0,  16,  17],
        [  0, -19,  20],
        [ 21,  22,   0]],

       [[ 24,   0,  26],
        [ 27,   0,  29],
        [  0, -31,   0],
        [ 33,  34,   0]],

       [[  0,  37,  38],
        [ 39,   0,  41],
        [ 42,  43,   0],
        [ 45,   0,  47]],

       [[ 49,  48,   0],
        [ 51,   0,  53],
        [  0,   0,  56],
        [  0,  59,  58]]])

输出数组中的掩码，其中 x 中的非零元素将被放置 -

In [97]: mask
Out[97]: 
array([[[ True,  True, False],
        [ True,  True, False],
        [ True,  True, False],
        [ True, False, False]],

       [[ True,  True, False],
        [ True,  True, False],
        [ True,  True, False],
        [ True,  True, False]],

       [[ True,  True, False],
        [ True,  True, False],
        [ True, False, False],
        [ True,  True, False]],

       [[ True,  True, False],
        [ True,  True, False],
        [ True,  True, False],
        [ True,  True, False]],

       [[ True,  True, False],
        [ True,  True, False],
        [ True, False, False],
        [ True,  True, False]]], dtype=bool)

最终输出-

In [98]: out
Out[98]: 
array([[[  2,   1,   0],
        [  3,   4,   0],
        [  8,   6,   0],
        [  9,   0,   0]],

       [[ 12, -14,   0],
        [ 16,  17,   0],
        [-19,  20,   0],
        [ 21,  22,   0]],

       [[ 24,  26,   0],
        [ 27,  29,   0],
        [-31,   0,   0],
        [ 33,  34,   0]],

       [[ 37,  38,   0],
        [ 39,  41,   0],
        [ 42,  43,   0],
        [ 45,  47,   0]],

       [[ 49,  48,   0],
        [ 51,  53,   0],
        [ 56,   0,   0],
        [ 59,  58,   0]]])

将零重新定位到多维 numpy 数组中最后一个维度的末尾

Relocate zeros to the end of the last dimension in multidimensional numpy array

python

arrays

numpy

vectorization

python-2.7