相乘的elementwise超过两个阵列
问题描述:
给出一个三维阵列和二维阵列的最终轴,相乘的elementwise超过两个阵列
a = np.arange(10*4*3).reshape((10,4,3))
b = np.arange(30).reshape((10,3))
如何运行跨过每个的最终轴的elementwise乘法,导致c
其中c
具有形状.shape
作为a
?即
c[0] = a[0] * b[0]
c[1] = a[1] * b[1]
# ...
c[i] = a[i] * b[i]
答
没有任何款项,减少参与,一个简单的broadcasting
将与np.newaxis/None
延伸b
到3D
后真正有效 -
a*b[:,None,:] # or simply a*b[:,None]
运行测试 -
In [531]: a = np.arange(10*4*3).reshape((10,4,3))
...: b = np.arange(30).reshape((10,3))
...:
In [532]: %timeit np.einsum('ijk,ik->ijk', a, b) #@Brad Solomon's soln
...: %timeit a*b[:,None]
...:
100000 loops, best of 3: 1.79 µs per loop
1000000 loops, best of 3: 1.66 µs per loop
In [525]: a = np.random.rand(100,100,100)
In [526]: b = np.random.rand(100,100)
In [527]: %timeit np.einsum('ijk,ik->ijk', a, b)
...: %timeit a*b[:,None]
...:
1000 loops, best of 3: 1.53 ms per loop
1000 loops, best of 3: 1.08 ms per loop
In [528]: a = np.random.rand(400,400,400)
In [529]: b = np.random.rand(400,400)
In [530]: %timeit np.einsum('ijk,ik->ijk', a, b)
...: %timeit a*b[:,None]
...:
10 loops, best of 3: 128 ms per loop
10 loops, best of 3: 94.8 ms per loop
答
使用np.einsum
:
c = np.einsum('ijk,ik->ijk', a, b)
快速检查:
print(np.allclose(c[0], a[0] * b[0]))
print(np.allclose(c[1], a[1] * b[1]))
print(np.allclose(c[-1], a[-1] * b[-1]))
# True
# True
# True
@BradSolomon这是正确的。增加了时间来确认这些。 – Divakar