How to make 2D jagged array using NumPy

Question

I want to make a 2D array that in the first row has 2 elements, in the second row has 4 elements, and in the third row has 6. Below is my code:

jagged_array = np.array([
[None, None], 
[None, None, None, None],
 [None, None, None, None, None, None]
])
print(jagged_array)
print(jagged_array.ndim)
print(jagged_array.shape)

However, the output doesn't look right. And the dimension is 1 (I expected 2). Here is the output:

[list([None, None]) list([None, None, None, None])
 list([None, None, None, None, None, None])]
1
(3,)

I want to know how I can make a 2D array with each row having different number of columns.

Answer 1

Based on this StackOverflow answer:

NumPy does not support jagged arrays natively. gives an array that may or may not behave as you expect.

A workaround using masked arrays can be as follows:

import numpy as np
import numpy.ma as ma

a = np.array([0, 1])
b = np.array([2, 3, 4, 5])
c = np.array([6, 7, 8, 9, 10, 11])

jagged_array = ma.vstack(
    [
        ma.array(np.resize(a, c.shape[0]), mask=[False, False, True, True, True, True]),
        ma.array(
            np.resize(b, c.shape[0]), mask=[False, False, False, False, True, True]
        ),
        c,
    ]
)
print(jagged_array)
print(jagged_array.ndim)
print(jagged_array.shape)

Your output would look like:

❯ python3 sample.py
[[0 1 -- -- -- --]
 [2 3 4 5 -- --]
 [6 7 8 9 10 11]]
2
(3, 6)