python - 将数组 (k,) 或 (k, n) 乘以一维数组 (k,)

Question

我有时会使用一维数组:

A = np.array([1, 2, 3, 4])

或 2D 阵列(使用 scipy.io.wavfile 读取单声道或立体声信号):

A = np.array([[1, 2], [3, 4], [5, 6], [7,8]])

不必用 if A.ndim == 2:... 来区分这两种情况，是否有一个简单的单行解决方案来乘以这个数组 A 通过一维数组 B = np.linspace(0., 1., 4)?

如果 A 是一维的，那么它就是 A * B，如果 A 是二维的，我的意思是将每一行乘以A 由 B 的每个元素组成。

---

注意:当使用 scipy.io.wavfile 读取单声道和立体声时，这个问题自然会出现。

Answer 1

方法 #1

我们可以使用einsum覆盖通用 ndarrays -

np.einsum('i...,i->i...',A,B)

这里的技巧是 ellipsis 广播第一个轴之后的尾随维度，因为它们在 A 中，并将它们保留在输出中，同时保持两个输入的第一个轴对齐，这是这里的预期乘法。将 A 作为 1D 数组，没有广播，本质上减少为:np.einsum('i,i->i',A,B) 引擎盖下。

示意图:

A   :  i x ....
B   :  i 
out :  i x ....

因此，覆盖具有任意维数的 A。

文档中有关使用 ellipsis 的更多信息:

To enable and control broadcasting, use an ellipsis. Default NumPy-style broadcasting is done by adding an ellipsis to the left of each term, like np.einsum('...ii->...i', a). To take the trace along the first and last axes, you can do np.einsum('i...i', a), or to do a matrix-matrix product with the left-most indices instead of rightmost, you can do np.einsum('ij...,jk...->ik...', a, b).

方法 #2

利用我们试图将 A 的第一个轴与一维数组 B 的唯一轴对齐的事实，我们可以简单地转置 A，与B相乘，最后转置回去-

(A.T*B).T