python - numpy数组的索引和？

Question

我想通过对另一个数组的索引值求和来生成一个 numpy 数组

例如，给定以下数组:

row_indices = np.array([[1, 1, 1], [0, 1, 1]])
col_indices = np.array([[0, 0, 1], [1, 1, 1]])
values = np.array([[2, 2, 3], [2, 4, 4]])

我想按以下方式设置一个新数组 indexed_sum:

for i in range(row_indices.size):
    indexed_sum[row_indices.flat[i], col_indices.flat[i]] += values.flat[i]

这样:

indexed_sum = np.array([[0, 2], [4, 11]])

但是，由于这是一个 Python 循环，并且这些数组可能非常大，因此需要花费大量时间。是否有一种有效的 numpy 方法可用于实现此目的？

Answer 1

您可能会通过 numba 获得成功，另一个 Python 包。我使用 %timeit 在 Jupyter notebook 中为以下两个函数计时。结果如下:

import numba
import numpy as np


# Your loop, but in a function.
def run_sum(row_indicies, col_indicies, values):
    indexed_sum = np.zeros((row_indicies.max() + 1, col_indicies.max() + 1))
    for i in range(row_indicies.size):
        indexed_sum[row_indicies.flat[i], col_indicies.flat[i]] += values.flat[i]
    return indexed_sum


# Your loop with a numba decorator.
@numba.jit(nopython=True)  # note you may be able to parallelize too
def run_sum_numba(row_indicies, col_indicies, values):
    indexed_sum = np.zeros((row_indicies.max() + 1, col_indicies.max() + 1))
    for i in range(row_indicies.size):
        indexed_sum[row_indicies.flat[i], col_indicies.flat[i]] += values.flat[i]
    return indexed_sum

我的示例数据有更大的意义:

row_id_big = np.random.randint(0, 100, size=(1000,))
col_id_big = np.random.randint(0, 100, size=(1000,))
values_big = np.random.randint(0, 10, size=(1000,))

结果:

%timeit run_sum(row_id_big, col_id_big, values_big)
# 1.04 ms ± 12.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit run_sum_numba(row_id_big, col_id_big, values_big)
# 3.85 µs ± 44.3 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

在此示例中，使用 numba 装饰器的循环要快几百倍。与您的示例相比，我对内存使用情况并不乐观。我必须初始化一个 numpy 数组，以便在某处放置数据，但如果您有更好的方法来执行该步骤，您可能能够进一步提高性能。

关于 numba 的注意事项:您需要运行一次循环才能开始看到主要的速度改进。您也许可以只使用像您在这里这样的玩具示例来初始化 jit，并看到相同的加速。