Python: How to split a list into ordered chunks

If I have the following list

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

Then

np.array_split([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], 3)

Returns

[array([0, 1, 2, 3]), array([4, 5, 6]), array([7, 8, 9])]

Is there a way to get the sub-arrays in the following order?

[array([0, 3, 6, 9]), array([1, 4, 7]), array([2, 5, 8])]

Solution 1:

As the lists are of differing lengths, a numpy.ndarray isn't possible without a bit of fiddling, as all sub-arrays must be the same length.

However, if a simple list meets your requirement, you can use:

l = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
l2 = []

for i in range(3):
    l2.append(l[i::3])

Output:

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

Or more concisely, giving the same output:

[l[i::3] for i in range(3)]

Solution 2:

Let's look into source code refactor of np.array_split:

def array_split(arr, Nsections):
    Neach_section, extras = divmod(len(arr), Nsections)
    section_sizes = ([0] + extras * [Neach_section + 1] + (Nsections - extras) * [Neach_section])
    div_points = np.array(section_sizes).cumsum()

    sub_arrs = []
    for i in range(Nsections):
        st = div_points[i]
        end = div_points[i + 1]
        sub_arrs.append(arr[st:end])

    return sub_arrs

Taking into account your example arr = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] and Nsections = 3 it will construct section sizes [0, 4, 3, 3] and dividing points [0, 4, 7, 10]. Then do something like this:

[arr[div_points[i]:div_points[i + 1]] for i in range(3)]

Trying to mimic behaviour of numpy, indeed,

def array_split_withswap(arr, N):
    sub_arrs = []
    for i in range(N):
        sub_arrs.append(arr[i::N])

Is the best option to go with (like in @S3DEV solution).