Numpy ordered subtraction of elements with full vectorization

Refresh

6 days ago

Views

4 time

0

Imagine a mxn array a, and a 1xn array b, we want to subtract b from a so that b is subtracted from the first element of a, then maximum of zero and b-a[0] is subtracted from a[1], and so on...

So:

x = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
a = np.repeat(x, 100000).reshape(10, 100000)
b = np.repeat(np.array([5]), 100000).reshape(1, 100000)

So we want to get: [ 0, 0, 1, 4, 5, 6, 7, 8, 9, 10], repeated 100,000 times.

I've managed the function below which provides the desired outcome:

def func(a, b):
    n = np.copy(a)
    m = np.copy(b)
    for i in range(len(n)):
        n[i] = np.where(n[i] >= m, n[i] - m, 0)
        m = np.maximum(0, m - a[i])
        if not m.any():
            return n
    return n

However, it's not fully vectorized. So:

>> timeit func(a, b)
3.23 ms ± 52.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Ideally, I'd like to get rid of the for-loop and make this as vectorized as possible.

Thanks.

0 answers