Question

I would like to extend a list while looping over it:

for idx in xrange(len(a_list)):
    item = a_list[idx]
    a_list.extend(fun(item))

(fun is a function that returns a list.)

Question:
Is this already the best way to do it, or is something nicer and more compact possible?

Remarks:

from matplotlib.cbook import flatten
a_list.extend(flatten(fun(item) for item in a_list))

should work but I do not want my code to depend on matplotlib.

for item in a_list:
    a_list.extend(fun(item))

would be nice enough for my taste but seems to cause an infinite loop.

Context:
I have have a large number of nodes (in a dict) and some of them are special because they are on the boundary.
‘a_list’ contains the keys of these special/boundary nodes. Sometimes nodes are added and then every new node that is on the boundary needs to be added to ‘a_list’. The new boundary nodes can be determined by the old boundary nodes (expresses here by ‘fun’) and every boundary node can add several new nodes.

Answer 1

Have you tried list comprehensions? This would work by creating a separate list in memory, then assigning it to your original list once the comprehension is complete. Basically its the same as your second example, but instead of importing a flattening function, it flattens it through stacked list comprehensions. [edit Matthias: changed + to +=]

a_list += [x for lst in [fun(item) for item in a_list] for x in lst]  

---

EDIT: To explain what going on.

So the first thing that will happen is this part in the middle of the above code:

[fun(item) for item in a_list]

This will apply fun to every item in a_list and add it to a new list. Problem is, because fun(item) returns a list, now we have a list of lists. So we run a second (stacked) list comprehension to loop through all the lists in our new list that we just created in the original comprehension:

for lst in [fun(item) for item in a_list]

This will allow us to loop through all the lists in order. So then:

[x for lst in [fun(item) for item in a_list] for x in lst]

This means take every x (that is, every item) in every lst (all the lists we created in our original comprehension) and add it to a new list.

Hope this is clearer. If not, I’m always willing to elaborate further.