有没有办法在两个二维数组的每一行上应用一个将两个一维数组作为参数的 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]])