首页 > 解决方案 > numpy广播在条件下乘以?

问题描述

我有两个数组,一个是 shape arr1.shape = (1000,2),另一个是 shape arr2.shape = (100,)

我想以某种方式乘以arr1[:,1]*arr2wherearr1[:,0] == arr2.index以获得最终形状arr_out.shape = (1000,)。arr1 的第一列本质上是一个id,其中以下条件成立:set(arr1[:,0]) == set(i for i in range(0,100)),即始终存在至少一个在 中arr2找到的值索引arr1[:,0]

我不太明白如何在 numpy 库中执行此操作,但觉得应该有一种使用numpy multiply的方法,如果有办法正确配置where条件?

我考虑过一个虚拟索引维度arr2可能会有所帮助吗?

可以使用以下代码片段生成一个玩具示例

arr2_length = 100
arr1_length = 1000

arr1 = np.column_stack(
    (np.random.randint(0,arr2_length,(arr1_length)),
     np.random.rand(arr1_length))
)
arr2 = np.random.rand(arr2_length)

# Doesn't work
arr2_b = np.column_stack((
    np.arange(arr2_length),
    np.random.rand(arr2_length)
))
# np.multiply(arr1[:,1],arr2_b[:,1], where=(arr1[:,0]==arr2_b[:,0]))

我有一种解决方案是利用 Pandas 中的左连接将较小的数组广播到相同长度的数组,然后相乘,如下所示:

df = pd.DataFrame(arr1).set_index(0).join(pd.DataFrame(arr2))
arr_out = (df[0]*df[1]).values

但我真的很想了解是否有一种原生的 numpy 方式来执行此操作,因为我觉得使用数据帧连接进行乘法不是一个非常易读的解决方案,并且可能会遭受过多的内存开销。

谢谢你的帮助。

标签: pythonarraysnumpy

解决方案


我相信这正是你想要的:

indices, values = arr1[:,0].astype(int), arr1[:,1]
arr_out = values * arr2[indices]

推荐阅读