2017-08-31 171 views
1

给出一个三维阵列和二维阵列的最终轴,相乘的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] 

回答

2

没有任何款项,减少参与,一个简单的broadcasting将与np.newaxis/None延伸b3D后真正有效 -

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 
+0

@BradSolomon这是正确的。增加了时间来确认这些。 – Divakar

1

使用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