如何加速python中的numpy數組布爾掩碼？

Question

我需要在 python 中加速以下布爾掩碼：

import numpy as np

# test dataset
n=1000000
mask = np.random.choice(a=[False, True], size=(n,), p=[0.8, 0.2])
arr = np.random.rand(n)

# the code I need to speed up:
res = arr[mask]

有什么想法可以讓它在python中更快嗎？

Answer 1

我寫了一個 numba 實現，它的性能與 numpy 幾乎相同。

import numba as nb
import numpy as np

@nb.njit(parallel=False, fastmath=True)
def array_masking_float(arr, mask, res):
    j=0
    for i in nb.prange(mask.shape[0]):
        if mask[i] == True:
            res[j] = arr[i]
            j += 1

誰能想出更好的解決方案？