如何用另一个数组中的值替换 numpy 矩阵列中的值？

Question

I have 2 numpy arrays:

a= np.array([[2, 1, 7],
             [7, 7, 3],
             [1, 7, 4]])


b= np.array([9,-1,17])

I would like to change the 7 in each column in a with the values from b , such that in the first column the 7 is replaced by the first value from b and in the second column the 7 s are replaced by the second value from b and so on.

My simple solution is:

for j in range(len(b)): 
    a[a[:,j]==7,j]=b[j]

array([[ 2,  1, 17],
       [ 9, -1,  3],
       [ 1, -1,  4]])

It works, but for very large matrices it is not fast enough. Is there another, faster way to do this?

Answer 1

Assuming the computed matrices are big, you can implement a fast parallel version using Numba . This implementation is much faster than the initial solution using pure-Python loops which create many small temporary arrays and an inefficient non-contiguous memory access pattern (eg. a[:,j] ). It is also significantly faster than using np.where(a == 7, b, a) due to the huge temporary arrays that needs to be filled and that may not fit in RAM (causing the OS to work with the very slow swap memory). Using multiple threads also provides a big speed up. Here is the code:

import numba as nb

@nb.njit('void(int_[:,::1], int_[::1])', parallel=True)
def compute(a, b):
    n, m = a.shape
    assert b.size == m
    for i in nb.prange(n):
        for j in range(m):
            if a[i,j] == 7:
                a[i,j] = b[j]

Here are results on my 6-core machine on a 100000x1000 matrix (with random 32-bit integers in 0..10):

For loop:  683 ms
np.where:  169 ms
Numba:      37 ms

This version is 18 times faster than the initial version and takes almost no more memory (it works in-place).