I have an one-dimensional array A, such that 0 <= A[i] <= 11, and I want to map A to an array B such that
for i in range(len(A)):
if 0 <= A[i] <= 2: B[i] = 0
elif 3 <= A[i] <= 5: B[i] = 1
elif 6 <= A[i] <= 8: B[i] = 2
elif 9 <= A[i] <= 11: B[i] = 3
How can implement this efficiently in numpy?
You need to use an int division by //3, and that is the most performant solution
A = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
B = A // 3
print(A) # [0 1 2 3 4 5 6 7 8 9 10 11]
print(B) # [0 0 0 1 1 1 2 2 2 3 3 3]
I would do something like dividing the values of the A[i] by 3 'cause you're sorting out them 3 by 3, 0-2 divided by 3 go answer 0, 3-5 go answer 1, 6-8 divided by 3 is equal to 2, and so on
I built a little schema here:
A[i] --> 0-2. divided by 3 = 0, what you wnat in array B[i] is 0, so it's ok
A[i] --> 3-5. divided by 3 = 1, and so on. Just use a method to make floor the value, so that it don't become float type.
Answers provided by others are valid, however I find this function from numpy quite elegant, plus it allows you to avoid for loop which could be quite inefficient for large arrays
import numpy as np
bins = [3, 5, 8, 9, 11]
B = np.digitize(A, bins)
//3
Something like this might work:
C = np.zeros(12, dtype=np.int)
C[3:6] = 1
C[6:9] = 2
C[9:12] = 3
B = C[A]
If you hope to expand this to a more complex example you can define a function with all your conditions:
def f(a):
if 0 <= a and a <= 2:
return 0
elif 3 <= a and a <= 5:
return 1
elif 6 <= a and a <= 8:
return 2
elif 9 <= a and a <= 11:
return 3
And call it on your array A:
A = np.array([0,1,5,7,8,9,10,10, 11])
B = np.array(list(map(f, A))) # array([0, 0, 1, 2, 2, 3, 3, 3, 3])
