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

[geremy condra]

On Mon, Nov 23, 2009 at 9:10 PM, geremy condra <debatem1 at gmail.com> wrote:
> On Mon, Nov 23, 2009 at 9:03 PM, geremy condra <debatem1 at gmail.com> wrote:
>> On Mon, Nov 23, 2009 at 7:05 PM, Paul Miller
>> <paul.w.miller.please.dont.spam.me at wmich.edu> wrote:
>>> I was wondering if there were any neat tools (like for instance,
>>> something from itertools) that would help me write the following function
>>> more elegantly.  The return value should, of course, be the complete $k$-
>>> partite graph $K_{n_1, n_2, \dots, n_k}$:
>>>
>>> def completeGraph (*ns):
>>>    '''
>>>    Returns the complete graph $K_{n_1, n_2, \dots, n_k}$ when passed
>>>    the sequence \code {n_1, n_2, \dots, n_k}.
>>>    '''
>>>    if len (ns) == 1:
>>>        return completeGraph ( * ([1] * ns[0]) )
>>>    n = sum (ns)
>>>    vertices = range (n)
>>>    partition_indices = [sum (ns[:i]) for i in range (len (ns))]
>>>    partite_sets = [vertices[partition_indices[i]:partition_indices[i+1]]
>>> \
>>>                    for i in range (len (partition_indices) - 1)]
>>>    partite_sets.append (vertices[partition_indices [-1]:] )
>>>
>>>    edges = []
>>>    for i in range (len (partite_sets)):
>>>        for j in range (i + 1, len (partite_sets)):
>>>            edges.extend ([ (u, v) for u in partite_sets [i] for v in \
>>>                           partite_sets [j] ])
>>>
>>>    return graph.Graph (vertices = vertices, edges = edges)
>>>
>>> Many thanks!
>>
>> Graphine does this with the following:
>>
>> from base import Graph
>>
>> def K(n):
>>        """Generates a completely connected undirected graph of size n.
>>
>>        The verticies are numbered [0, n).
>>
>>        The edges are named after the verticies they connect such that
>>        an edge connected verticies 1 and 2 is named (1,2).
>>        """
>>        # create the graph
>>        k = Graph()
>>        # generate all the nodes
>>        for i in range(n):
>>                k.add_node(i)
>>        # generate all the edges
>>        for i in range(n):
>>                for j in range(i+1, n):
>>                        k.add_edge(i, j, (i,j), is_directed=False)
>>        # return the graph
>>        return k
>>
>>
>> Disclaimer: I'm the author of graphine.
>>
>> Geremy Condra
>>

Alright, how does this look:

def k_partite(*partition_sizes):
	g = Graph()
	for pos, num_nodes in enumerate(partition_sizes):
		for i in range(num_nodes):
			n = g.add_node(name=(pos,i), partition=pos)
	for node1 in g.nodes:
		for node2 in g.nodes:
			if node1.partition != node2.partition:
				g.add_edge(node1, node2, is_directed=False)
	return g

Geremy Condra