Sticky cumsum/cumprod numpy

QuantumPotatoe :

Foreword: I'll probably update the title as I suspect I only lack the correct wording for my problem. (Which makes it a potential duplicate, sorry about that). I'd appreciate suggestions to do so.

Here's functionally what I want to accomplish:

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

To try and be clearer, here's a naive implementation (Goal being obviously to do that with numpy methods):

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

Edit 1: As pointed out in the comments, [0, 0] in first position has an undefined behaviour. That was a simplified case of my actual context (which would be irrelevant). It's assumed that [0, 0] will never happen in first position.

Quang Hoang :

You can use np.select to mask (1,0) and (0,1). Then use this answer to fill in the nan with previous values:

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]

Output:

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

Guess you like

Origin http://43.154.161.224:23101/article/api/json?id=194064&siteId=1