如何为Numpy二维数组每行元素按行号递增对应数值

我需要实现给NumPy二维数组的每行元素分别加上对应数值：第0行所有元素加1，第1行所有元素加2，以此类推。

已掌握的一维数组实现方式 #

对于一维数组，我可以直接通过以下代码实现需求：

import numpy as np

y = np.random.random(size=(10,))
y += np.arange(start=1, stop=11, step=1)

运行结果示例：

[ 1.66704813  2.90625487  3.961399    4.64969495  5.5286593   6.31736865
  7.24056032  8.25632077  9.02458071 10.1816335 ]

二维数组遇到的问题 #

但处理形状为(10,3)的二维数组时，直接套用上述写法无法生效：

y = np.random.random(size=(10,3))
y += np.arange(start=1, stop=11, step=1)  # 无法生效，需要让每行的所有列都加上对应行号+1的数值

期望得到的结果示例：

[[1.66704813 1.90625487 1.961399  ]
 [2.64969495 2.5286593  2.31736865]
 [3.24056032 3.25632077 3.02458071]
 [4.1816335  4.13044436 4.13969493]
 [5.48485214 5.58980863 5.59890548]
 [6.58126238 6.30004674 6.730429  ]
 [7.76121817 7.94366992 7.1039714 ]
 [8.81874117 8.00239219 8.16807975]
 [9.64509574 9.75334071 9.59641831]
 [10.11176155 10.71027095 10.77104173]]

解决方案 #

问题出在NumPy的广播机制上，一维数组(10,)和二维数组(10,3)的形状不兼容，无法直接广播。我们需要把一维数组转换为(10,1)的形状，这样就能匹配二维数组的行数，进而广播到每行的所有列。

有两种常用的转换方式：

1. 使用reshape调整维度：

import numpy as np

y = np.random.random(size=(10,3))
# 将arange生成的数组转为(10,1)形状，实现广播
y += np.arange(1, 11).reshape(10, 1)

2. 使用[:, np.newaxis]增加维度：

import numpy as np

y = np.random.random(size=(10,3))
y += np.arange(1, 11)[:, np.newaxis]

两种方式都能让每行对应的数值（第i行对应i+1）被广播到该行的所有列，最终实现每行元素都加上对应数值的需求，得到和示例一致的结果。

内容的提问来源于stack exchange，提问作者Tim