Python Numpy - 创建长度基于一维数组的二维数组

Question

抱歉让标题困惑，但不确定如何让它更简洁。这是我的要求:

arr1 = np.array([3,5,9,1])
arr2 = ?(arr1)

arr2 将是:

[
[0,1,2,0,0,0,0,0,0],
[0,1,2,3,4,0,0,0,0],
[0,1,2,3,4,5,6,7,8],
[0,0,0,0,0,0,0,0,0]
]

它不需要根据最大值而变化，形状是预先知道的。所以开始我已经能够得到零的形状:

arr2 = np.zeros((len(arr1),max_len))

然后我当然可以像这样在 arr1 上做一个 for 循环:

for i, element in enumerate(arr1):
    arr2[i,0:element] = np.arange(element)

但这可能需要很长时间，而且这里的两个维度都相当大(arr1 是几百万行，max_len 大约是 500)。在 numpy 中是否有一种干净优化的方法来执行此操作？

Answer 1

基于 @Divakar 几年前发布的“填充”想法:

In [161]: res = np.arange(9)[None,:].repeat(4,0)
In [162]: res[res>=arr1[:,None]] = 0
In [163]: res
Out[163]: 
array([[0, 1, 2, 0, 0, 0, 0, 0, 0],
       [0, 1, 2, 3, 4, 0, 0, 0, 0],
       [0, 1, 2, 3, 4, 5, 6, 7, 8],
       [0, 0, 0, 0, 0, 0, 0, 0, 0]])