CODEWITHSUNDEEP
python
Matthias 009
Matthias 009

Reputation: 1712

append/extend list in loop

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.

Upvotes: 2

Views: 12422

Answers (4)

Lucas Vazquez
Lucas Vazquez

Reputation: 1882

Using generator

original_list = [1, 2]
original_list.extend((x for x in original_list[:]))

# [1, 2, 1, 2]

Upvotes: 1

Jochen Ritzel
Jochen Ritzel

Reputation: 107608

As you want to extend the list, but loop only over the original list, you can loop over a copy instead of the original:

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

Upvotes: 1

Hod - Monica's Army
Hod - Monica's Army

Reputation: 875

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.

Upvotes: 2

phihag
phihag

Reputation: 287835

Using itertools, it can be written as:

import itertools
a_list += itertools.chain(* itertools.imap(fun, a_list))

or, if you're aiming for code golf:

a_list += sum(map(fun, a_list), [])

Alternatively, just write it out:

new_elements = map(fun, a_list) # itertools.imap in Python 2.x
for ne in new_elements:
  a_list.extend(ne)

Upvotes: 2

Related Questions