How can I assign values from one array to another according to the index more efficiently?

Question

I am trying to replace the values of one array with another according to how many ones are in the source array. I assign a value from a given index in the replacement array based on the sum. Thus, if there are 2 ones in a row, it assigns a value of l1[1] to the species, and if there is one unit, it assigns a value of l1[0] to the output.

It will be better seen in a specific example:

import numpy as np

l1 = np.array([4, 5])
x112 = np.array([[0, 0], [0, 1], [1, 1], [0, 0], [1, 0], [1, 1]])

array([[0, 0],
       [1, 0],
       [1, 1],
       [0, 0],
       [1, 0],
       [1, 1]])

Required output:

[[0]
 [4]
 [5]
 [0]
 [4]
 [5]]

I did this by counting the units in each row and assigning accordingly using np.where :

x1x2 = np.array([0, 1, 2, 0 1, 2]) #count value 1
x1x2 = np.where(x1x2 != 1, x1x2, l1[0]) 
x1x2 = np.where(x1x2 != 2, x1x2, l1[1])             
print(x1x2)

output

[0 4 5 0 4 5]

Could this be done more effectively?

Answer 1

Okay I actually gave devectorizing your code a shot. First the vectorized NumPy you have:

def op(x112, l1):
    # bit of cheating, adding instead of counting 1s
    x1x2 = x112[:,0] + x112[:,1]

    x1x2=np.where(x1x2 != 1, x1x2, l1[0])
    x1x2=np.where(x1x2 != 2, x1x2, l1[1])
    return x1x2

The most efficient alternative is to loop through x112 only once, so let's do a Numba loop.

import numba as nb

@nb.njit
def loop(x112, l1):
    d0, d1 = x112.shape
    x1x2 = np.zeros(d0, dtype = x112.dtype)
    for i in range(d0):
        # actually count the 1s
        num1s = 0
        for j in range(d1):
            if x112[i,j] == 1:
                num1s += 1
        
        if num1s == 1:
            x1x2[i] = l1[0]
        elif num1s == 2:
            x1x2[i] = l1[1]
    return x1x2

Numba loop has a ~9-10x speed improvement on my laptop.

%timeit op(x112, l1)
8.05 µs ± 34.4 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit loop(x112, l1)
873 ns ± 5.09 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

As @Mad_Physicist requested, timings with a bigger array. I'm including his advanced-indexing method too.

x112 = np.random.randint(0, 2, size = (100000, 2))
l1_v2 = np.array([0,4,5])

%timeit op(x112, l1)
1.35 ms ± 27.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit loop(x112, l1)
956 µs ± 2.78 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit l1_v2[x112.sum(1)]
1.2 ms ± 1.05 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

EDIT: Okay maybe take these timings with a grain of salt because when I went to restart the IPython kernel and reran this stuff, op(x112, l1) improved to 390 µs ± 22.1 µs per loop while the other methods retained the same performance (971 µs, 1.23 ms).

Answer 2

You can use direct indexing:

l1 = np.array([0, 4, 5])
x112 = np.array([[0, 0], [0, 1], [1, 1], [0, 0], [1, 0], [1, 1]])

result = l1[x112.sum(1)]

This works if you're at liberty to prepend the zero to l1 at creation time. If not:

result = np.r_[0, l1][x112.sum(1)]