python - 使用 numpy 数组加速函数

Question

from numba import jit

@jit
def interim_mk(x, unique_x):
    """

    :param x:
    :param unique_x:
    :return:
    """
    tp = np.zeros(unique_x.shape)

    for i in range(len(unique_x)):
        tp[i] = sum(x == unique_x[i])

    return tp

在上面的函数中，我使用 jit 来尝试加快速度。然而，这似乎没有帮助。 x 和 unique_x 都是 numpy 数组，有没有办法加快计算速度(不使用 cython)

Answer 1

对于正数的情况

您可以使用np.bincount -

count = np.bincount(x)
out = count[count!=0]

使用unique_x -

out = np.bincount(x)[unique_x]

对于一般情况

out = np.bincount(np.searchsorted(unique_x, x))

<小时/>

当然，我们可以直接从 np.unique 调用中获取计数，如果这就是我们 unique_x 的方式 -

out = np.unique(x, return_counts=1)[1]