python - 如何使用条件对 numpy 3D 数组中的值求和？

Question

我有一个包含 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 工具上搜索，但没有找到任何与我的具体问题相关的内容。如果有人可以向我指出正确的文档/工具，我将不胜感激。谢谢。

Answer 1

这是一种方法 -

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])]

另一种更多地使用 bool 数组来提高效率 -

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