用于列表的numpy范围（矢量化计算）

Question

我想从数组b创建一个二维矩阵a，其中a包含每个矩阵列的range_stop值。

例如，我想使用a = [2,3]

b = [[0, 0],
     [1, 1],
     [2, 2],
     [NaN, 3]]

（向量化计算）最有效的方法是什么？我当前的代码是：

a = [2,3]
b = np.zeros((max(a)+1,len(a)))
b.fill(np.nan)
for i,ai in enumerate(a):
    b[:ai, i] = np.arange(ai)

Answer 1

您可以先使用arange创建二维repeat

a = np.asarray([2, 3])
b = np.repeat(np.arange(np.max(a) + 1, dtype=float)[:, None], len(a), axis=1)
# array([[0., 0.],
#        [1., 1.],
#        [2., 2.],
#        [3., 3.]])

，然后将每列与a进行比较以填写np.nan s

b[b > a] = np.nan
# array([[ 0.,  0.],
#        [ 1.,  1.],
#        [ 2.,  2.],
#        [nan,  3.]])