Any elegant way to construct the complete $k$-partite graph in Python?

[Malte Helmert]

Paul Miller wrote:
> On Mon, 23 Nov 2009 19:57:05 -0800, Richard Thomas wrote:
> 
>> Not sure exactly how you're representing graphs, this seems like the
>> simplest way of listing the edges.
>>
>> def complete_partite(*sizes):
>>     total = sum(sizes)
>>     nodes, edges = range(total), []
>>     for group in xrange(len(sizes)):
>>         low = sum(sizes[:group-1])
>>         high = sum(sizes[:group])

I think this has a conceptual off-by-one error. Add

           print group, low, high

to see what I mean (especially the first iteration). It still works, but
I think this would be clearer:

           low = sum(sizes[:group])
           high = sum(sizes[:group + 1])

or to avoid doing essentially the same summation twice:

           low = sum(sizes[:group])
           high = low + sizes[group]

>>         edges.extend((i, j) for i in xrange(low, high)
>>                             for j in xrange(high, total))
>>     return nodes, edges

Here's a variant that uses a running total instead of recomputing the
sum in every iteration, thus getting rid of xrange(len(...)).

def complete_partite(*sizes):
    total = sum(sizes)
    nodes, edges = range(total), []
    curr_total = 0
    for size in sizes:
        edges.extend((i, j) for i in xrange(curr_total, curr_total+size)
                            for j in xrange(curr_total+size, total))
        curr_total += size
    return nodes, edges

Finally, here is a variant that is a bit shorter because it produces the
edges in a different way and hence gets rid of the need for knowing the
total up front and uses total as running total instead. It has the
drawback of not generating the edges in ascending order though, so I
think the previous one is nicer:

def complete_partite(*sizes):
    total, edges = 0, []
    for size in sizes:
        edges.extend((i, j) for i in xrange(total)
                            for j in xrange(total, total + size))
        total += size
    return range(total), edges

Finally, here's a variation on the same theme:

def complete_partite(*sizes):
    nodes, edges = [], []
    for size in sizes:
        partition = xrange(len(nodes), len(nodes) + size)
        edges.extend((i, j) for i in nodes for j in partition)
        nodes.extend(partition)
    return nodes, edges

Malte