有没有一种方法可以应用numpy函数，将两个1d数组作为两个2d数组中每一行的参数？

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]])