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

Question

我正在尝试运行类似以下的内容：

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

其中array1和array2都是形状为（m，n）的二维numpy数组。我的目标是使用数组1和数组2的每一行分别运行n次np.bincount()。

我尝试使用np.apply_along_axis()，但据我所知，这仅允许在数组1的每一行上运行该函数，而没有使用数组2的每一行作为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)]

- Kronos

---

总是有助于给出输入和预期输出的小例子... - Julien

apply_along... 不是一个性能工具，即使在它工作的情况下也不是。它不会编译函数并在 c 中执行循环。 - hpaulj

没有通用的方法来做到这一点。如果每对行都必须调用一次函数，无论迭代方法或包装器如何，都需要时间。有时可以重新思考问题并一次解决所有行的问题。但是，bincount通过在1d中工作来提高速度。 - hpaulj

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

- Jérôme Richard

---

这是我的当前解决方案。自从提出问题后，我意识到我的输出大小不同，现在使用minlength参数来确保输出大小相同。我将更新我的示例以反映这一点。 - Kronos