Numpy trim_zeros в 2D или 3D - PullRequest
Новая
       13

Numpy trim_zeros в 2D или 3D

0 голосов
/ 30 апреля 2019

Как удалить начальные / конечные нули из массива NumPy? Trim_zeros работает только для 1D.

1 Ответ

0 голосов
/ 30 апреля 2019

Следующая функция работает для любого измерения: 

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)

, которое можно проверить с помощью: 

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