Algorithm for distributing a number between certain number of chunks

Question

I am looking for an algorithm to split up a number between a certain number of the chunk. The real example is: assigning initial capacity to all the workers based on the total number of tasks.

eg If there are 3 persons and 6 tasks. Each will have an initial capacity = 6 / 3 = 2

I come up with an implementation. But, I am not sure it is an optimal way to achieve it.

def capacity_distribution(task_size, people_size):
    """
    It distributes initial capacity to each person.
    """
    div = task_size // people_size

    remainder = task_size % people_size

    first_chunk_size = people_size - remainder
    second_chunk_size = people_size - first_chunk_size

    first_capacity_list = []
    second_capacity_list = []
    first_capacity_list = [div for _ in range(first_chunk_size)]
    second_capacity_list = [div + 1 for _ in range(second_chunk_size)]
    first_capacity_list.extend(second_capacity_list)

    return first_capacity_list



print(capacity_distribution(6, 2))
print(capacity_distribution(7, 3))
print(capacity_distribution(11, 3))
print(capacity_distribution(18, 5))

Output :

[3, 3]
[2, 2, 3]
[3, 4, 4]
[3, 3, 4, 4, 4]

Is there any other efficient way for this one?

Answer 1

Maybe:

def capacity_distribution(task_size, people_size):
    each = math.floor(task_size/people_size) # everyone gets this many
    extra = task_size % people_size # this many get 1 extra

    distribution = [each for x in range(people_size)]
    for x in range(people_size):
        if x < extra:
            distribution[x] += 1

    return distribution

Answer 2

How about this?

def capacity_distribution(task_size, people_size):
    modulus = task_size % people_size
    div = task_size // people_size

    dist = []

    if modulus:
        dist.append(div)
        dist.extend(capacity_distribution(task_size-div, people_size-1))
    else:
        dist.extend([div]*people_size)

    return dist

Answer 3

My solution would be:

def cap_distr(ts, ps):
    l = [ts // ps for _ in range(ps)]
    for i in range(ts % ps):
        l[i] += 1
    return l

test:

for test in [(6, 2), (7, 3), (11, 3), (18, 5)]: print(cap_distr(*test))
[3, 3]
[3, 2, 2]
[4, 4, 3]
[4, 4, 4, 3, 3]

PS: if you don't mind importing numpy, it would be even one line shorter:

from numpy import ones

def cap_distr(ts, ps):
    arr = ones(ps) * (ts // ps)
    arr[:ts % ps] += 1
    return arr.astype(int)