CODEWITHSUNDEEP
pythonlistgenerator
Scriptim
Scriptim

Reputation: 1832

Generator for splitting list (sequence) in N chunks in Python

for i in generate_chunks([1, 2, 3, 4, 5], 2):
    print(i)

# [[1], [2, 3, 4, 5]]
# [[1, 2], [3, 4, 5]]
# [[1, 2, 3], [4, 5]]
# [[1, 2, 3, 4], [5]]


for i in generate_chunks([1, 2, 3, 4, 5], 3):
    print(i)

# [[1], [2], [3, 4, 5]]
# [[1, 2], [3], [4, 5]]
# [[1, 2, 3], [4], [5]]
# [[1, 2], [3, 4], [5]]
# ...

How can I implement generate_chunks(list, n)?

Essentially what generate_chunks does is splitting list in n chunks and yield a list of these chunks.

For clarification, n refers to the number of chunks, not the length of the chunks.

The order in which those lists of chunks are yielded is irrelevant, however the order of the elements in the initial list is important, so that for a given list [1, 2, 3] the result [[1], [2, 3]] would be valid whereas [[1], [3, 2] would be invalid.

(Preferably without using a 3rd-party library)

Upvotes: 1

Views: 398

Answers (1)

John Coleman
John Coleman

Reputation: 51998

Here is an itertools based approach:

import itertools

def chunks(items, cutpoints):
    return [items[i:j] for i,j in zip([0] + cutpoints, cutpoints + [len(items)])]

def generate_chunks(items, n):
    indices = range(1,len(items))
    for cutpoints in itertools.combinations(indices,n-1):
        yield chunks(items,list(cutpoints))

For example:

>>> for c in generate_chunks([1,2,3,4,5],4): print(c)

[[1], [2], [3], [4, 5]]
[[1], [2], [3, 4], [5]]
[[1], [2, 3], [4], [5]]
[[1, 2], [3], [4], [5]]

Upvotes: 7

Related Questions