如何对两个2d numpy数组对应行调用接收双1d数组参数的numpy函数

实现方案 #

你可以通过给每行索引增加偏移量的方式，单次调用np.bincount实现逐行加权计数，完全规避Python层循环，性能最优：

核心逻辑

给array1的每行元素加上对应行的偏移值，让不同行的bin索引互不重叠，全局计算bincount后再reshape回(m, minlength)的目标形状。

示例代码

import numpy as np

# 输入示例
array1 = np.array([[1,2,3],[4,5,6]])
array2 = np.array([[7,8,9],[10,11,12]])
minlength = 7
m = array1.shape[0]

# 计算每行的索引偏移量
row_offset = np.arange(m)[:, np.newaxis] * minlength
# 给array1的每行加上对应偏移，避免不同行索引重叠
array1_with_offset = array1 + row_offset

# 拉平后单次计算bincount，再reshape为目标形状
result = np.bincount(
    array1_with_offset.ravel(),
    weights=array2.ravel(),
    minlength=m * minlength
).reshape(m, minlength)

输出验证

上述代码运行后result的输出为：

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

和你预期的逐行调用np.bincount的结果完全一致。

注：你示例中写的np.bincounts为笔误，numpy原生对应方法为np.bincount。

注意事项

• 该方案要求array1中的元素值均小于你指定的minlength，和直接调用np.bincount的输入要求一致。如果array1存在超过minlength的元素，可先将minlength调整为array1.max() + 1再执行上述逻辑。
• 该方案为完全向量化实现，当行数m很大时，性能远高于逐行循环或apply_along_axis实现。

内容的提问来源于stack exchange，提问作者Vhagar