How to subset a numpy array of different lengths

Question

I have a numpy array and would like to subset the first two arrays of each element in an ndarray.

Here is an example array:

import numpy as np

a1 = np.array([[ 1,  2,  3],
               [ 4,  5,  6]])

a2 = np.array([[ 7,  8,  9],
               [10, 11, 12],
               [13, 14, 15],
               [16, 17, 18]])

a3 = np.array([[19, 20, 21],
               [22, 23, 24],
               [25, 26, 27]])

A = np.array([a1, a2, a3])

print("A =\n", A)

Which prints:

A =
   [array([[ 1,  2,  3],
           [ 4,  5,  6]])
    array([[ 7,  8,  9],
           [10, 11, 12],
           [13, 14, 15],
           [16, 17, 18]])
    array([[19, 20, 21],
           [22, 23, 24],
           [25, 26, 27]])]

The desired result is as follows:

A =
   [array([[ 1,  2,  3],
           [ 4,  5,  6]])
    array([[ 7,  8,  9],
           [10, 11, 12]])
    array([[19, 20, 21],
           [22, 23, 24]])]

To print the equivalent object, you could do

print(np.array([a1[0:2], a2[0:2], a3[0:2]]))

But I want to directly get what is desired using A.

What is the correct way of doing this in numpy?

Edit: I would like to subset the array without looping. Alternative ways of structuring the arrays so that they can be directly indexed are okay too. Any numpy function to avoid looping is fair game.

Answer 1

a = [i[0:2] for i in A]

This will work!