Как суммировать значения в массивном трехмерном массиве с условными? - PullRequest
Новая
       14

Как суммировать значения в массивном трехмерном массиве с условными?

1 голос
/ 29 октября 2019

У меня есть массив Numpy 3d со значениями 0 и 1, что-то вроде этого: 

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? Я искал на форумах нудистские инструменты, но я не нашел ничего относительно моей конкретной проблемы. Если бы кто-то мог указать мне на правильную документацию / инструмент, я был бы признателен. Спасибо.

1 Ответ

2 голосов
/ 29 октября 2019

Вот один из способов - 

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
...