python - 有没有办法应用一个 numpy 函数，该函数将两个一维数组作为两个二维数组的每一行的参数？

Question

我正在尝试运行类似的东西:

 np.bincount(array1, weights = array2, minlength=7)

其中 array1 和 array2 都是形状为 (m,n) 的 2d n numpy 数组。我期望的目标是 np.bincount() 对 array1 和 array2 的每一行运行 n 次

我试过使用 np.apply_along_axis()但据我所知，这只允许函数在 array1 的每一行上运行，而不使用 array2 的每一行作为 np.bincount 的参数。我希望找到一种方法来使用 numpy 函数而不是迭代来干净地完成此操作，因为这是一个性能关键函数，但到目前为止找不到另一种方法。

例如，给定这些数组:

array1 = [[1,2,3],[4,5,6]]
array2  = [[7,8,9],[10,11,12]]

我想计算:

[np.bincounts([1,2,3], weights = [7,8,9],minlength=7),  np.bincounts([4,5,6], weights = [10,11,12], minlength=7)]

Answer 1

一个简单的解决方案就是使用综合列表:

result = [np.bincount(v, weights=w) for v,w in zip(array1, array2)]

因为生成的数组可以有不同的大小(实际上在您的示例中确实有不同的大小)，所以结果不能是 Numpy 数组，而是常规列表。大多数 Numpy 函数无法处理可变大小数组 的列表，甚至无法生成它们。

如果数组中有很多行，您可以使用 Numba 的 JIT(在这种情况下最终使用 Cython)来降低 CPython 解释器循环的成本。请注意，为了提高性能，必须在调用 Numba 函数之前将输入数组转换为 Numpy 数组。如果您知道所有数组的大小都相同，则可以使用 Numba 编写更高效的实现(通过预分配生成的数组并自己进行 bincount)。

---

更新

对于固定大小的数组，这是 Numba 中的一个快速实现:

import numpy as np
import numba as nb

array1 = np.array([[1,2,3],[4,5,6]], dtype=np.int32)
array2  = np.array([[7,8,9],[10,11,12]], dtype=np.int32)

@nb.njit('i4[:,::1](i4[:,::1],i4[:,::1])')
def compute(array1, array2):
    assert array1.shape == array2.shape
    n, m = array1.shape
    res = np.zeros((n, 7), dtype=np.int32)
    for i in range(n):
        for j in range(m):
            v = array1[i, j]
            assert v>=0 and v<7  # Can be removed if the input is safe
            res[i, v] += array2[i, j]
    return res

result = compute(array1, array2)

# result is
# array([[ 0,  7,  8,  9,  0,  0,  0],
#       [ 0,  0,  0,  0, 10, 11, 12]])