Numba 用数组快速替换键与值的方法

Question

我想用具有重复元素的大尺寸array中的values替换keys。我正在尝试numba和numpy映射方法。两种方法的代码如下。你知道吗

import numpy as np
from numba import njit, prange

array1 = np.arange(150*150*150, dtype=int)
array2 = np.arange(150*150*150, dtype=int)
array = np.concatenate((array1, array2))

keys = np.arange(50)
values = -1 * np.arange(50)

## Numba Approach
@njit(parallel=True)
def numba_replace(array, keys, values):

    for i in prange(len(keys)):
        for j in prange(len(array)):
            if array[j] == keys[i]:
               array[j] = values[i]


## numpy approach
def numpy_replace(array, keys, values):

    mapp = np.arange(array.size)
    mapp[keys] = values
    mapped = mapp[array]

    return mapped

## Performance 
%%timeit
numba_replace(array, keys, values)
# 117 ms ± 969 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%%timeit
numpy_replace(array, keys, values)
# 61.2 ms ± 159 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

虽然numpy_replace比numba_replace快2倍，但我不喜欢使用它，因为我的数组大小非常大(3000 x 3000 x 3000)，numpy方法会创建一个new array来提高内存使用率。有什么方法可以让numba\u替换更快，或者有什么方法在这个过程中不会创建新的数组？你知道吗

Answer 1

改进Numba方法（降低复杂性）

由于您只想更改相对较少的值，因此可以使用集合来确定是否必须更改实际数组元素。 此外，您还可以使用search\u sorted来获得正确的键、值对。对于这个小例子，差异并不是很大，但是如果问题的规模增加，差异会变得更大。你知道吗

实施

import numpy as np
from numba import njit, prange

@njit(parallel=True)
def numba_replace(array, keys, values):
    ind_sort=np.argsort(keys)
    keys_sorted=keys[ind_sort]
    values_sorted=values[ind_sort]
    s_keys=set(keys)

    for j in prange(array.shape[0]):
        if array[j] in s_keys:
            ind = np.searchsorted(keys_sorted,array[j])
            array[j]=values_sorted[ind]
    return array

计时

import numpy as np
from numba import njit, prange

array1 = np.arange(150*150*150, dtype=int)
array2 = np.arange(150*150*150, dtype=int)
array = np.concatenate((array1, array2))

#to get proper timings do nothing here
#changing the array in-place will obviously have 
#an influence on the timings, because there are no values to change in the second run
keys = np.arange(50)
values = np.arange(50)

%timeit numba_replace(array, keys, values)
# 20.1 ms ± 1.95 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit numpy_replace(array, keys, values)
# 51.3 ms ± 392 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Answer 2

我猜：

array[keys] = values

在numpy中执行此操作，而不创建任何新数组

编辑：检查命令是否执行与numpy_replace函数相同的操作：

mapped = numpy_replace(array, keys, values)
array[keys] = values
print(all(mapped == array)) #  > True