[FIXED] How to vectorize such an algorithm in python?

Issue

I have four (nx1) dimensional arrays named a, b, c, and F. I want to run this algorithm without any loops.

for i in range(n):
    if a[i] < b[i]:
        F[i] = 0
    elif a[i] > c[i]:
        F[i] = 1
    elif b[i] <= a[i] <= c[i]:
        F[i] = 2

I want to write this code in a more vectorized way to make it more efficient as my dataset is quite large.

Solution

If you care about performance, why not try numba? It might get 10X faster than logical operations while saving memory at the same time. As a bonus, the loop code you wrote will be kept intact, only through an @njit decorator in front of the function.

from numba import njit

@njit
def loop(F, a, b, c):
  for i in range(F.shape[0]):
    if a[i] < b[i]:
      F[i] = 0
    elif a[i] < c[i]:
      F[i] = 1
    elif b[i] <= a[i] <= c[i]:
      F[i] = 2

Compare with vectorized solution by @NiziL using sizes of 100 and 1000 vectors,

timeit(lambda: loop(F, a, b, c))
timeit(lambda: idx(F, a, b, c))

Gives:

# 1.0355658 (Size: 100, @njit loop)
# 4.6863165 (Size: 100, idx)

# 1.9563843 (Size: 1000, @njit loop)
# 16.658198 (Size: 1000, idx)

Answered By - AboAmmar