Python 如何在numpy 3D数组中使用条件值求和？

我有一个带有0和1值的Numpy 3d数组，类似于：

array([[[ 1, 1, 0, 1],
        [ 0, 0, 1, 1]],

       [[ 1, 1, 1, 1],
        [ 0, 1, 0, 1]]])

我想在数组的“行”中“添加”（+运算）每个值，具体条件如下：如果我有连续的“1”值，我将求和。如果我有0，我会保持原样。值为“0”后，我重新开始计数

我希望得到的输出是：

array([[[ 2, 0, 1],
        [ 0, 0, 2]],

       [[ 4],
        [ 0, 1, 0, 1]]])

输出可以是不同大小的“线”。我还能和numpy一起做吗？ 我在论坛上搜索了numpy工具，但没有找到任何与我的具体问题相关的东西。如果有人能告诉我正确的文档/工具，我将不胜感激。谢谢。

这里有一种方法-

def sum_groups(a):
    z = np.zeros(a.shape[:-1] + (1,), dtype=a.dtype)
    b = np.concatenate((z,a,z),axis=-1)

    c = b.ravel()
    count = np.diff(np.flatnonzero(c[:-1]!=c[1:]))

    m2 = c[1:]>c[:-1]
    c[1:][m2] = count[::2]

    m3 = c==0
    m3[1:][m2] = 1
    m4 = m3.reshape(b.shape)
    m4[...,0] = 0
    m4[...,-1] = 0
    v = c[m3]

    rc = m4.sum(2).ravel()
    out = np.split(v,rc[:-1].cumsum())
    return out

样本运行-

In [7]: a
Out[7]: 
array([[[1, 1, 0, 1],
        [0, 0, 1, 1]],

       [[1, 1, 1, 1],
        [0, 1, 0, 1]]])

In [8]: sum_groups(a)
Out[8]: [array([2, 0, 1]), array([0, 0, 2]), array([4]), array([0, 1, 0, 1])]

另一个是为了提高效率而更多地使用布尔数组-

def sum_groups_v2(a):
    p1 = a==1
    z1 = np.zeros(a.shape[:-1] + (1,), dtype=bool)
    b1 = np.concatenate((z1,p1,z1),axis=-1)

    c1 = b1.ravel()
    count1 = np.diff(np.flatnonzero(c1[:-1]!=c1[1:]))
    m33 = np.r_[False,c1[:-1]<c1[1:]].reshape(b1.shape)
    pp = np.zeros(b1.shape, dtype=int)
    pp[m33] = count1[::2]
    m33[...,1:-1][~p1] = 1
    v2 = pp[m33]
    rc2 = m33.sum(2).ravel()
    out2 = np.split(v2,rc2[:-1].cumsum())
    return out2

---

非常感谢你！它使用的高级函数（与我的基本编程级别相比）是我自己在短时间内找不到的！：）

from scipy.ndimage import label

def sum_groups_v3(a):
    z = np.zeros(a.shape[:-1] + (1,), dtype=a.dtype)
    b = np.concatenate((z,a),axis=-1)

    c = b.ravel()

    l = label(c)[0]
    unq,idxs,count = np.unique(l,return_index=True, return_counts=True)

    l[c==1] = -1
    l[idxs] = count
    p = l.reshape(b.shape)[...,1:]
    out = [j[j>=0] for i in p for j in i]
    return out