Масштабировать каждый столбец островов до их длины в двумерном массиве NumPy

Question

У меня есть такой массив значений, например:

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

, и я хотел бы найти последовательные пиксели со значением 1 в каждом столбце и установить для этих значений пикселей полученную длину, чтобы получитьэтот вывод:

array([[0,  0,  0,  3,  0,  4],
       [0,  0,  0,  3,  0,  4],
       [2,  1,  1,  3,  0,  4],
       [2,  0,  0,  0,  2,  4],
       [0,  0,  0,  0,  2,  0],
       [1,  1,  0,  0,  0,  1]])

Спасибо за вашу помощь

Answer 1

Подход № 1

def scaleby_grouplen(ar):
    a = ar==1
    a1 = np.pad(a, ((1, 1), (0, 0)), 'constant')
    a2 = a1.ravel('F')
    idx = np.flatnonzero(a2[1:] != a2[:-1])
    start, stop = idx[::2], idx[1::2]
    id_ar = np.zeros(len(a2), dtype=int)
    id_ar[start+1] = 1
    idx_ar = id_ar.cumsum()-1
    lens = stop - start
    out = a*lens[idx_ar].reshape(-1,a.shape[0]+2).T[1:-1]
    return out

Подход № 2

В качестве альтернативы можно использовать np.maximum.accumulate для замены cumsum детали-

def scaleby_grouplen_v2(ar):
    a = ar==1
    a1 = np.pad(a, ((1, 1), (0, 0)), 'constant')
    a2 = a1.ravel('F')
    idx = np.flatnonzero(a2[1:] != a2[:-1])
    start, stop = idx[::2], idx[1::2]
    id_ar = np.zeros(len(a2), dtype=int)
    id_ar[start+1] = np.arange(len(start))
    idx_ar = np.maximum.accumulate(id_ar)
    lens = stop - start
    out = a*lens[idx_ar].reshape(-1,a.shape[0]+2).T[1:-1]
    return out

Подход № 3

Использование np.repeat для повторения длин групп и, следовательно, заполнения -

def scaleby_grouplen_v3(ar):
    a = ar==1
    a1 = np.pad(a, ((1, 1), (0, 0)), 'constant')
    a2 = a1.ravel('F')
    idx = np.flatnonzero(a2[1:] != a2[:-1])
    lens = idx[1::2] - idx[::2]
    out = ar.copy()
    out.T[a.T] = np.repeat(lens, lens)
    return out

Пробный прогон -

In [177]: a
Out[177]: 
array([[0, 0, 0, 1, 0, 1],
       [0, 0, 0, 1, 0, 1],
       [1, 1, 1, 1, 0, 1],
       [1, 0, 0, 0, 1, 1],
       [0, 0, 0, 0, 1, 0],
       [1, 1, 0, 0, 0, 1]])

In [178]: scaleby_grouplen(a)
Out[178]: 
array([[0, 0, 0, 3, 0, 4],
       [0, 0, 0, 3, 0, 4],
       [2, 1, 1, 3, 0, 4],
       [2, 0, 0, 0, 2, 4],
       [0, 0, 0, 0, 2, 0],
       [1, 1, 0, 0, 0, 1]])

Бенчмаркинг

Другой подход (ы) -

from numpy import array
from itertools import chain, groupby

# @timgeb's soln
def chain_groupby(a):
    groups = (groupby(col, bool) for col in a.T)
    unrolled = ([(one, list(sub)) for one, sub in grp] for grp in groups)
    mult  = ([[x*len(sub) for x in sub] if one else sub for one, sub in grp] for grp in unrolled)
    chained = [list(chain(*sub)) for sub in mult]
    result = array(chained).T
    return result

Сроки -

In [280]: np.random.seed(0)

In [281]: a = np.random.randint(0,2,(1000,1000))

In [282]: %timeit chain_groupby(a)
1 loop, best of 3: 667 ms per loop

In [283]: %timeit scaleby_grouplen(a)
100 loops, best of 3: 17.7 ms per loop

In [284]: %timeit scaleby_grouplen_v2(a)
100 loops, best of 3: 17.1 ms per loop

In [331]: %timeit scaleby_grouplen_v3(a)
100 loops, best of 3: 18.6 ms per loop

Answer 2

Я чувствую, что может быть более короткое решение pandas, если вы хотите использовать эту библиотеку, но на данный момент это лучшее, что я мог придумать.

>>> from numpy import array
>>> from itertools import chain, groupby
>>> 
>>> a = array([[0,  0,  0,  1,  0,  1],
...        [0,  0,  0,  1,  0,  1],
...        [1,  1,  1,  1,  0,  1],
...        [1,  0,  0,  0,  1,  1],
...        [0,  0,  0,  0,  1,  0],
...        [1,  1,  0,  0,  0,  1]])
>>> 
>>> groups = (groupby(col, bool) for col in a.T)
>>> unrolled = ([(one, list(sub)) for one, sub in grp] for grp in groups)
>>> mult  = ([[x*len(sub) for x in sub] if one else sub for one, sub in grp] for grp in unrolled)
>>> chained = [list(chain(*sub)) for sub in mult]
>>> result = array(chained).T
>>> 
>>> result
array([[0, 0, 0, 3, 0, 4],
       [0, 0, 0, 3, 0, 4],
       [2, 1, 1, 3, 0, 4],
       [2, 0, 0, 0, 2, 4],
       [0, 0, 0, 0, 2, 0],
       [1, 1, 0, 0, 0, 1]])