How to create a chunked list of lists with almost-equal lengths?

I need to split a number into 3 parts. The following function (source) does that correctly:

def parts(num: int, div: int) -> list[int]:
    """Split a number into equal parts."""
    return [num // div + (1 if x < num % div else 0) for x in range(div)]

So

assert parts(8, 3) == [3, 3, 2]
assert parts(9, 3) == [3, 3, 3]
assert parts(10, 3) == [4, 3, 3]

From that list, I need a list of lists containing a range from 0 to num + 3, such that every sublist

• at an even index has a length of 1
• at an odd index has a length of n, where n is fetched from parts

assert expected(8, 3) == [[0], [1, 2, 3], [4], [5, 6, 7], [8], [9, 10]]
assert expected(9, 3) == [[0], [1, 2, 3], [4], [5, 6, 7], [8], [9, 10, 11]]
assert expected(10, 3) == [[0], [1, 2, 3, 4], [5], [6, 7, 8], [9], [10, 11, 12]]

This is what I’ve tried:

def actual(num: int, div: int) -> list[list[int]]:
    matrix: list[list[int]] = []
    for i, length in enumerate(parts(num, div)):
        base = i * (length + 1)
        matrix.append([base])
        matrix.append([base + j + 1 for j in range(length)])

    return matrix

However, it only works when num is divisible by div. How can I fix it?

These are the outputs for my function:

assert actual(8, 3) == [[0], [1, 2, 3], [4], [5, 6, 7], [6], [7, 8]]
assert actual(9, 3) == [[0], [1, 2, 3], [4], [5, 6, 7], [8], [9, 10, 11]]
assert actual(10, 3) == [[0], [1, 2, 3, 4], [4], [5, 6, 7], [8], [9, 10, 11]]

>Solution :

It looks like each sub list contains incrementing numbers. This is what itertools.count can give you.

I’ve adapted your function to use it like this:

from itertools import count

def actual(num: int, div: int) -> list[list[int]]:
    matrix: list[list[int]] = []
    base = count(0)
    for length in parts(num, div):
        matrix.append([next(base)])
        matrix.append([next(base) for _ in range(length)])

    return matrix