重置为 0 的多维 Numpy 数组的累积和

Question

我有一个多维 NumPy 数组 (16x212)，我想计算 212 列中每一列的累积和 - 但是，如果有0 之间。

例如array([0, 1, -1, 1, 0, -1, 1, 0, 0, 1, 0, -1, 0, 0, 1, -1], ...)

我观看了以下视频，其中您将 0 和 1 的数组累积在一起，但是如果存在负数，该方法将不再起作用。

有人可以帮帮我吗？

Answer 1

fill_zeros_with_last2d 基于this answer。

这里的方法是取沿轴的累积总和，然后从后面的列中减去累积到最后一个零的总和。

import numpy as np

def fill_zeros_with_last2d(adj, arr ):
    row, col = np.indices( adj.shape )
    col[(adj == 0) & (arr != 0) ] = 0
    col = np.maximum.accumulate(col, axis = 1 )
    # Find the last column with a value

    return adj[row, col]

def cum_reset_2d( arr ):
    accum = arr.cumsum( axis = 1 )
    adj = accum * ( arr == 0 )
    return accum - fill_zeros_with_last2d( adj, arr )

在一个小数组上测试

np.random.seed( 1235 )

arr = np.random.randint( -1, 2, size = ( 5, 10 ))
arr
""" 
array([[ 1,  1,  1,  0, -1, -1, -1,  1,  1, -1],
       [ 1,  1, -1,  1,  1,  1, -1,  0, -1,  1],
       [ 0, -1,  1,  0,  0, -1,  1,  1,  0, -1],
       [ 1,  0, -1,  0, -1,  1,  0,  0,  1,  1],
       [-1,  0,  0, -1, -1,  1,  1, -1,  1, -1]])
"""

cum_reset_2d( arr )
""" 
array([[ 1,  2,  3,  0, -1, -2, -3, -2, -1, -2],
       [ 1,  2,  1,  2,  3,  4,  3,  0, -1,  0],
       [ 0, -1,  0,  0,  0, -1,  0,  1,  0, -1],
       [ 1,  0, -1,  0, -1,  0,  0,  0,  1,  2],
       [-1,  0,  0, -1, -2, -1,  0, -1,  0, -1]])
"""

注意我无法让它工作，因此需要在使用前进行测试。原始版本不适用于 arr[3]，但它们现在看起来都还可以。可能还存在一些问题。我的 adj 值为零，必须保留但正在被覆盖。