Search code examples
pythonlistnumpysplitbinning

Is there another way to split a list into bins of equal size and put the remainder if any into the first bin?


Given a sorted list:

x = list(range(20))

I could split the list into equal sizes and put the remainder into the left bins as such:

def split_qustions_into_levels(questions, num_bins=3):
    num_questions = len(questions)
    equal_size = int(num_questions / num_bins)
    slices = [equal_size] * num_bins
    slices[0] += len(questions) % num_bins
    return [[questions.pop(0) for _ in questions[:s]] for s in slices]

If I have 3 bins from the list of 20 items, I should get a output list of list with sizes (7,7,6):

>>> x = list(range(20))
>>> split_qustions_into_levels(x, 3)
[[0, 1, 2, 3, 4, 5, 6, 7], [8, 9, 10, 11, 12, 13], [14, 15, 16, 17, 18, 19]]

If I want 6 bins from the list of 20 items, I should get the output list of list with size (5,3,3,3,3,3):

>>> x = list(range(20))
>>> split_qustions_into_levels(x, 6)
[[0, 1, 2, 3, 4],
 [5, 6, 7],
 [8, 9, 10],
 [11, 12, 13],
 [14, 15, 16],
 [17, 18, 19]]

Is there another way to do this without the messy calculation of slice and equal sizes and the left popping each item into a list of list?

Is there a numpy solution?


Solution

  • Maybe array_split() is what you want.

    a = np.arange(20)
    np.array_split(a, 6)