粘孜然

Question

前言：我可能会更新标题，因为我怀疑我只是缺少解决问题的正确措辞。（这可能会造成重复，对此表示抱歉）。我很感谢你的建议

以下是我想要实现的功能：

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

使用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)

Answer 2

您可以使用np.select来屏蔽(1,0)和(0,1)。然后使用this answer用前面的值填充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.])