Python Numpy - Create 2d array where length is based on 1D array

Question

Sorry for confusing title, but not sure how to make it more concise. Here's my requirements:

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

arr2 would then be:

[
[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]
]

It doesn't need to vary based on the max, the shape is known in advance. So to start I've been able to get a shape of zeros:

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

And then of course I could do a for loop over arr1 like this:

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

but that would likely take a long time and both dimensions here are rather large (arr1 is a few million rows, max_len is around 500). Is there a clean optimized way to do this in numpy?

Answer 1

Building on a 'padding' idea posted by @Divakar some years ago:

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]])

Answer 2

Try this with itertools.zip_longest -

import numpy as np
import itertools

l = map(range, arr1)
arr2 = np.column_stack((itertools.zip_longest(*l, fillvalue=0)))
print(arr2)

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]])