Group a numpy array

Question

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 ?

Answer 1

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]

Answer 2

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.

Answer 3

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)

Answer 4

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]

Answer 5

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])