如何在有条件的numpy 3D数组中求和?

时间:2019-10-29 13:04:17

标签: python numpy numpy-ndarray

我有一个带有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工具,但没有发现有关我的特定问题的任何信息。如果有人可以指出合适的文档/工具,我将不胜感激。谢谢。

1 个答案:

答案 0 :(得分:2)

这是一种方法-

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

还有一个region-based labelling-

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