Липкий сперма / комар numpy

Question

Предисловие: Я, вероятно, обновлю название, так как подозреваю, что мне не хватает только правильной формулировки для моей проблемы. (Что делает его потенциальным дубликатом, извините за это). Буду признателен за предложения сделать это.

Вот функционально то, что я хочу выполнить sh:

import numpy as np

a = np.array([1, 0, 0, 0, 1, 0])
b = np.array([0, 0, 1, 0, 0, 0])

# What I want:
result = [1, 1, 0, 0, 1, 1]
# [1, 0] -> 1 
# [0, 1] -> 0
# [0, 0] -> the last value

Чтобы попытаться быть более понятным, вот наивная реализация (цель, очевидно, состоит в том, чтобы сделать это с помощью numpy методов):

def sticky_cumsum(a, b):
    result = []
    for val_a, val_b in zip(a, b):
        if val_a:
            results.append(1)
        elif val_b: # assuming both can't be true
            results.append(0)
        else: # assuming first iteration won't be [0, 0]
            results.append(results[-1])
    return results

Редактировать 1: Как указано в комментариях, [0, 0] в первой позиции имеет неопределенное поведение. Это был упрощенный случай моего фактического контекста (который был бы неуместен). Предполагается, что [0, 0] никогда не произойдет в первой позиции.

Answer 1

Вы можете использовать np.select для маскировки (1,0) и (0,1). Затем используйте этот ответ , чтобы заполнить nan предыдущими значениями:

arr = np.select((a==1, b==1), (1,0), np.nan)

# inspired by the linked question
mask = np.isnan(arr)
idx = np.where(~mask,np.arange(len(mask)), 0)
np.maximum.accumulate(idx,out=idx)
out = arr[idx]

Вывод:

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

Answer 2

Альтернатива с использованием Cython

Если вы можете использовать Cython, вы можете использовать что-то вроде следующего, которое можно запустить из ноутбука Jupyter с помощью следующих

%load_ext Cython

Тогда в отдельной ячейке запустите.

%%cython -a

from cython cimport boundscheck, wraparound
cimport numpy as np
import numpy as np

@boundscheck(False)
@wraparound(False)
def cython_sticky_cumsum(int[::1] a, int[::1] b):

    cdef size_t i, N = a.size
    cdef np.ndarray[np.int32_t, ndim=1] result = np.empty(N, dtype=np.int32)

    for i in range(N):
        if a[i] == 1:
            result[i] = 1
        elif b[i] == :
            result[i] = 0
        else:
            result[i] = result[i-1]

    return result

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

a = np.array([1, 0, 0, 0, 1, 0])
b = np.array([0, 0, 1, 0, 0, 0])

cython_sticky_cumsum(a, b)
# array([1, 1, 0, 0, 1, 1])

Необработанный тест

Для гораздо больших массивов, таких как

a = np.tile(np.array([1, 0, 0, 0, 1, 0]), 1000000)
b = np.tile(np.array([0, 0, 1, 0, 0, 0]), 1000000)

и выполнение тестов,

%timeit cython_sticky_cumsum(a,b)
%timeit sticky_cumsum(a, b)

выходы

28.4 ms ± 1.86 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
2.5 s ± 97.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)