How to do vectorization using numpy on the condition that if larger then overwrite?-CodePudding

I need to write something like below. But how to use np.where or any other vectorization method to speed it up?

import numpy as np
A = np.arange(5)
B = [4, 3 ,1 ,2]
C = np.zeros(A.shape)
for a, b in zip(A, B):
    if C[b] < a - b:
        C[b] = a - b
print(C)

CodePudding user response：

Warning: if B contains duplicate elements, this is very hard (if even possible) to vectorize the loop since vectorized Numpy operation operate on the whole array every time and do not consider aliasing. The following solution (like the one of @Ali_Sh) assume there are no duplicate. It also assume C is zeros as in the example and the type of A and B is the same.

Without duplicates

A (twice) faster approach than the one of @Ali_Sh is to first compute the difference and use a multiplication instead of np.where which is slower in this case (this is because np.where tends to use slow conditional branches internally). Here is the code:

C = np.zeros(A.shape)
tmp = A[:len(B)] - B
C[B] = (tmp > 0) * tmp

Note that using B = np.fromiter(B, dtype=A.dtype) make the code faster (as other ones) since Numpy reconvert B from a list to an array every time and np.array is not the most efficient way to do that.

With duplicates

If there are duplicates and aliasing can happen, then you can use Numba to speed up the operation. Other Numpy solutions can give wrong results as opposed to this one. Here is the code:

import numba as nb

@nb.njit
def compute_with_aliasing(A, B):
    C = np.zeros(A.shape, dtype=A.dtype)
    for a, b in zip(A, B):
        if C[b] < a - b:
            C[b] = a - b
    return C

compute_with_aliasing(A, np.fromiter(B, dtype=A.dtype))

Note that the first use of the Numba function will be slower due to the compilation time. This approach should also be faster than the previous one in practice.

CodePudding user response：

You can do this as:

diff = A[:len(B)] - B                 # [-4 -2  1  1]
cond = np.less(C[B], diff)            # [False False  True  True]
C[B] = np.where(cond, diff, C[B])     # [0. 1. 1. 0. 0.]

A solution was written in the comments by @Szczesny, which was much like this one and can be rewritten with using len(B) instead -1 as:

C[B] = np.where(A[:len(B)] - B > 0, A[:len(B)] - B, 0)

It might be deleted due to similarity to my answer. I put it here again because it was one-line code and will remove if Szczesny put it as a new answer.