2D或3D中的Numpy trim_zeros [英] Numpy trim_zeros in 2D or 3D

问题描述

如何从NumPy数组中删除前导/尾随零? Trim_zeros 仅适用于1D.

How to remove leading / trailing zeros from a NumPy array? Trim_zeros works only for 1D.

推荐答案

以下功能适用于任何尺寸:

The following function works for any dimension:

def trim_zeros(arr, margin=0):
    '''
    Trim the leading and trailing zeros from a N-D array.

    :param arr: numpy array
    :param margin: how many zeros to leave as a margin
    :returns: trimmed array
    :returns: slice object
    '''
    s = []
    for dim in range(arr.ndim):
        start = 0
        end = -1
        slice_ = [slice(None)]*arr.ndim

        go = True
        while go:
            slice_[dim] = start
            go = not np.any(arr[tuple(slice_)])
            start += 1
        start = max(start-1-margin, 0)

        go = True
        while go:
            slice_[dim] = end
            go = not np.any(arr[tuple(slice_)])
            end -= 1
        end = arr.shape[dim] + min(-1, end+1+margin) + 1

        s.append(slice(start,end))
    return arr[tuple(s)], tuple(s)

可以通过以下方式进行测试:

Which can be tested with:

test = np.zeros((3,4,5,6))
test[1,2,2,5] = 1
trim_zeros(test, margin=1)

这篇关于2D或3D中的Numpy trim_zeros的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！