如何在Python的列表列表中对元素进行分组？

Question

I'm trying to efficientlly group / nest a list of vertices by their neighbourhood size, into a list of lists of vertices.

The neighbourhood size is a property of a vertex v can be obtained by calling len(v.neighbours) .

The input I have is an unsorted list of vertices. The output I'm trying to obtain should look like this:

[[all vertices with len(v.neighbours) == 1], [... == 2], [... == 4]]    

It should be a list of lists where each sublist contains vertices with the same neighbourhood size, sorted from small to big, without empty lists. I don't need the indices of the sublists to map to the neighbourhood size of the contained vertices.

I know how to achieve this with list comprehension, but it's rather inefficient:

def _group(V: List[Vertex], max: int) -> List[List[Vertex]]:
    return [[v for v in V if v.label == i] for i in range(max)]

Additionally, I don't want to pass the maximum neighbourhood size as a parameter, but calculate that during the grouping, and I'm looking for a way to filter out empty lists during the grouping as well.

I've looked into more efficient ways to group the vertices, for example by using a dictionary as intermediary step, but I haven't managed to produce a working result.

Could anyone tell me the most efficient way to group / nest the list of vertices?

Thanks in advance, and sorry if this has been posted before, but I couldn't find what I was looking for in another question.

Answer 1

One pass over the input, put the result in an intermediate dictionary, work the dictionary into the output you want.

temp_result = defaultdict(list)

for v in vertices:
    temp_result[neighborhood_size(v)].append(v)

max_size = max(temp_result.keys())

return_val = list()
for i in range(max_size):
    if temp_result[i]: # check if empty
        return_val.append(temp_result[i])

Answer 2

You can build it this way:

from collections import defaultdict

# Create a dict {nb_of_neighbours:[corresponding vertices]}
vertices_by_neighbours = defaultdict(list)
for v in vertices:
    vertices_by_neighbours[len(v.neighbours)].append(v)

# Create the output list, sorted by number of neighbours    
out = []
for nb_neighbours in sorted(vertices_by_neighbours):
    out.append(vertices_by_neighbours[nb_neighbours])

# Max number of neighbours, in case you want it...
max_neighbours = max(vertices_by_neighbours)