Home > OS >  Numpy split unequally
Numpy split unequally

Time:08-25

I have this array:

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

I've been trying to make something like this:

Split this array into 3 (unequally) and I want it to return something like this:

  1. 1st element -> 1st array,
  2. 2nd element -> 2nd array,
  3. 3rd element -> 3rd array,
  4. 4th element -> 1st array,
  5. etc.
[1, 4, 7, 10], [2, 5, 8], [3, 6, 9]

I tried something like this

import numpy as np
arr = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])

new_arrays = np.array_split(arr, 3)
print(new_arrays)

but It splits it like this:

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

I tried hsplit, It kinda works, but requires equal division, which is not possible for my dataset.

CodePudding user response:

I hope I've understood your question right. You can create n new lists and then use itertools.cycle to append values to each sublists:

from itertools import cycle

def my_split(arr, n):
    out = [[] for _ in range(n)]
    for v, l in zip(arr, cycle(out)):
        l.append(v)
    return out


lst = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

new_arrays = my_split(lst, 3)
print(new_arrays)

Prints:

[[1, 4, 7, 10], [2, 5, 8], [3, 6, 9]]
  • Related