如何沿某个轴广播分配numpy数组?

时间:2019-04-24 03:41:49

标签: python arrays numpy

我可以“广播”-分配带有列的矩阵吗?也就是说,我希望所有第一行都被分配给该列中的第一个条目,所有第二行都被分配给该列中的第二个条目,依此类推。

示例

s = np.arange(15).reshape(3, 5)
"""
s
Out[18]: 
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14]])
"""
s[[0, 1], 1:] = np.array([1, 2])

抛出:

Traceback (most recent call last):
  File "D:\Anaconda37\lib\site-packages\IPython\core\interactiveshell.py", line 2961, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-19-606ac1668601>", line 1, in <module>
    s[[0, 1], 1:] = np.array([1, 2])
ValueError: shape mismatch: value array of shape (2,) could not be broadcast to indexing result of shape (2,4)

有没有办法以广播的方式做到这一点?

1 个答案:

答案 0 :(得分:1)

IIUC

s[[0,1],1:]=np.array([1, 2])[:,None]#numpy broadcast here 
s
Out[1169]: 
array([[ 0,  1,  1,  1,  1],
       [ 5,  2,  2,  2,  2],
       [10, 11, 12, 13, 14]])