Generate all equal-sized partitions of a set

Question

I need a generator that gets as input a set of 'agents' and a set of 'items', and generates all partitions in which each agent gets the same number of items. For example:

>>> for p in equalPartitions(["A","B"], [1,2,3,4]): print(p)
{'A': [1, 2], 'B': [3, 4]}
{'A': [1, 3], 'B': [2, 4]}
{'A': [1, 4], 'B': [2, 3]}
{'A': [2, 3], 'B': [1, 4]}
{'A': [2, 4], 'B': [1, 3]}
{'A': [3, 4], 'B': [1, 2]}

For two agents this is easy (assuming the number of items is even):

itemsPerAgent = len(items) // len(agents)
for bundle0 in itertools.combinations(items, itemsPerAgent):
        bundle1 =  [item for item in items if item not in bundle0]
        yield {
            agents[0]: list(bundle0),
            agents[1]: bundle1
            }

For three agents this becomes more complicated:

itemsPerAgent = len(items) // len(agents)
for bundle0 in itertools.combinations(items, itemsPerAgent):
    bundle12 =  [item for item in items if item not in bundle0]
    for bundle1 in itertools.combinations(bundle12, itemsPerAgent):
        bundle2 =  [item for item in bundle12 if item not in bundle1]
        yield {
            agents[0]: list(bundle0),
            agents[1]: list(bundle1),
            agents[2]: bundle2
            }

Is there a more general solution, that works for any number of agents?

Answer 1

Here's a recursive solution, which works by allocating the appropriate number of items to the first agent, and handing the rest of the problem off to further invocations of itself:

from itertools import combinations

def part(agents, items):
    if len(agents) == 1:
        yield {agents[0]: items}
    else:
        quota = len(items) // len(agents)
        for indexes in combinations(range(len(items)), quota):
            remainder = items[:]
            selection = [remainder.pop(i) for i in reversed(indexes)][::-1]
            for result in part(agents[1:], remainder):
                result[agents[0]] = selection
                yield result

In the trivial case of a single agent, we yield a single dictionary and terminate.

If there's more than one agent, we:

1. Generate all combinations of indexes into items that should be allocated to the first agent.

2. Pop the items at those indexes from a copy of items in reverse order (to avoid messing up the indexes) into selection, reversing the result back again afterwards with [::-1] to maintain the expected order.

3. Call part() recursively on the remaining agents and items.

4. Add the selection we already made to each result yielded by those recursive calls, and yield that.

Here it is in action:

>>> for p in part(["A", "B"], [1, 2, 3, 4]):
...     print(p)
... 
{'A': [1, 2], 'B': [3, 4]}
{'A': [1, 3], 'B': [2, 4]}
{'A': [1, 4], 'B': [2, 3]}
{'A': [2, 3], 'B': [1, 4]}
{'A': [2, 4], 'B': [1, 3]}
{'A': [3, 4], 'B': [1, 2]}

>>> for p in part(["A", "B", "C"], [1, 2, 3, 4, 5, 6, 7, 8, 9]):
...     print(p)
... 
{'A': [1, 2, 3], 'B': [4, 5, 6], 'C': [7, 8, 9]}
{'A': [1, 2, 3], 'B': [4, 5, 7], 'C': [6, 8, 9]}
{'A': [1, 2, 3], 'B': [4, 5, 8], 'C': [6, 7, 9]}
{'A': [1, 2, 3], 'B': [4, 5, 9], 'C': [6, 7, 8]}
{'A': [1, 2, 3], 'B': [4, 6, 7], 'C': [5, 8, 9]}
  # [...]    
{'A': [7, 8, 9], 'B': [3, 4, 5], 'C': [1, 2, 6]}
{'A': [7, 8, 9], 'B': [3, 4, 6], 'C': [1, 2, 5]}
{'A': [7, 8, 9], 'B': [3, 5, 6], 'C': [1, 2, 4]}
{'A': [7, 8, 9], 'B': [4, 5, 6], 'C': [1, 2, 3]}

>>> for p in part(["A", "B", "C"], [1, 2, 3, 4, 5, 6, 7]):
...     print(p)
... 
{'A': [1, 2], 'B': [3, 4], 'C': [5, 6, 7]}
{'A': [1, 2], 'B': [3, 5], 'C': [4, 6, 7]}
{'A': [1, 2], 'B': [3, 6], 'C': [4, 5, 7]}
{'A': [1, 2], 'B': [3, 7], 'C': [4, 5, 6]}
  # [...]
{'A': [6, 7], 'B': [2, 5], 'C': [1, 3, 4]}
{'A': [6, 7], 'B': [3, 4], 'C': [1, 2, 5]}
{'A': [6, 7], 'B': [3, 5], 'C': [1, 2, 4]}
{'A': [6, 7], 'B': [4, 5], 'C': [1, 2, 3]}

As you can see, it handles cases where items can't be equally divided between agents. In addition, unlike the various permutations()-based answers, it doesn't waste work computing duplicate results, so runs a lot faster than them.

Answer 2

If you have a permutations function which can handle repeated elements in the input without producing duplicated permutations in the output, you can do this pretty efficiently. Unfortunately, itertools.permutations doesn't do what we want (len(list(itertools.permutations('aaa'))) is 6, not 1, which we want).

Here's a permutation function I wrote for some previous question, which happens to do the right thing with duplicated input values:

def permutations(seq):
    perm = sorted(seq) # the first permutation is the sequence in sorted order
    while True:
        yield perm

        # find largest index i such that perm[i] < perm[i+1]
        for i in range(len(perm)-2, -1, -1):
            if perm[i] < perm[i+1]:
                break
        else: # if none was found, we've already found the last permutation
            return

        # find the largest index j such that perm[i] < perm[j] (always exists)
        for j in range(len(perm)-1, -1, -1):
            if perm[i] < perm[j]:
                break

        # Swap values at indexes i and j, then reverse the values from i+1
        # to the end. I've packed that all into one operation, with slices.
        perm = perm[:i]+perm[j:j+1]+perm[-1:j:-1]+perm[i:i+1]+perm[j-1:i:-1]

Now, here's how to use it to assign items to your agents:

def equal_partitions(agents, items):
    items_per_agent, extra_items = divmod(len(items), len(agents))
    item_assignments = agents * items_per_agent + agents[:extra_items]
    for assignment in permutations(item_assignments):
        result = {}
        for agent, item in zip(assignment, items):
            result.setdefault(agent, []).append(item)
        yield result

The first lines build a list of references to our agents that is the same length as the list of items. Each agent is repeated as many times as the number of items they will receive. If the items list can't be divided exactly evenly, I add some extra references to the first few agents. You could add something else if you prefer (such as ['extra'] * extra_items, perhaps).

The main loop runs on the permutations of the assignments list. It then runs an inner loop that matches an agent from the permutation to the corresponding item, and packs up the result into a dictionary in the format you want.

This code should be asymptotically optimal in both time and space for any number of agents or items. That said, it may still be slow, since it relies on my permutation function written in pure Python rather than a faster implementation in C. It may be that there's an efficient way to get the permutations we want using itertools, but I'm not exactly sure how.

Answer 3

A terribly memory-inefficient solution, but quite short and more "pythonic". Also, the order of dictionaries in the result is quite arbitrary, imo.

import itertools as it
from pprint import pprint
from time import time

agents = ('a', 'b', 'c')
items = (1,2,3,4,5,6,7,8,9)

items_per_agent = int(len(items)/len(agents))

def split_list(alist,max_size=1):
    """Yield successive n-sized chunks from alist."""
    for i in range(0, len(alist), max_size):
        yield alist[i:i+max_size]

def my_solution():
    # I have put this into one-liner below
    # combos = set()
    # i=0
    # for perm in it.permutations(items, len(items)):
    #   combo = tuple(tuple(sorted(chunk)) for chunk in split_list(perm, max_size=items_per_agent))
    #   combos.add(combo)
    #   print(combo, i)
    #   i+=1

    combos = {tuple(tuple(sorted(chunk)) for chunk in split_list(perm, max_size=items_per_agent)) for perm in it.permutations(items, len(items))}

    # I have put this into one-liner below
    # result = []
    # for combo in combos:
    #   result.append(dict(zip(agents,combo)))

    result = [dict(zip(agents,combo)) for combo in combos]

    pprint(result)

my_solution()

Answer 4

# -*- coding: utf-8 -*-

from itertools import combinations
from copy import copy


def main(agents, items):
    if len(items) % len(agents):
        return []

    result = [{'remain': items}]

    part_size = len(items) / len(agents)

    while True:
        for item in result[:]:
            remain_agent = set(agents) - set(item.keys())
            if not remain_agent:
                continue

            result.remove(item)

            agent = remain_agent.pop()

            for combination in combinations(item['remain'], part_size):
                current_item = copy(item)
                current_item.update({agent: combination, 'remain': list(set(item['remain']) - set(combination))})
                result.append(current_item)

            break
        else: 
            break

    for item in result:
        item.pop('remain', None)
    return result


if __name__ == '__main__':
    agents = ['A', 'B', 'C']
    items = [1, 2, 3, 4, 5, 6]

    t = main(agents, items)

Here it is in action:

In [3]: agents = ['A', 'B']

In [4]: items = [1, 2, 3, 4]

In [5]: result = main(agents, items)

In [6]: for item in result:
   ...:     print item
   ...:
{'A': (1, 2), 'B': (3, 4)}
{'A': (1, 3), 'B': (2, 4)}
{'A': (1, 4), 'B': (2, 3)}
{'A': (2, 3), 'B': (1, 4)}
{'A': (2, 4), 'B': (1, 3)}
{'A': (3, 4), 'B': (1, 2)}

Answer 5

from itertools import combinations,permutations
def get(items, no_of_agents):
    def chunks(l, n):
        """Yield successive n chunks from l."""
        rt = []
        ln = len(l) // n
        for i in range(0, len(l) - ln - 1, ln):
            rt.append(l[i:i + ln])
        rt.append(l[i + ln:])
        return rt

    for i in permutations(items, len(items)):
        yield chunks(i,no_of_agents)

def get_equal_partitions(items, agents):
    for i in get(items, len(agents)):
        yield dict(zip(agents, i))

items = [i for i in range(4)]
agents = ["A","B","C"]

for i in get_equal_partitions(items,agents):
    print(i)