How can I, in python, iterate over multiple 2d lists at once, cleanly?

Question:

If I’m making a simple grid based game, for example, I might have a few 2d lists. One might be for terrain, another might be for objects, etc. Unfortunately, when I need to iterate over the lists and have the contents of a square in one list affect part of another list, I have to do something like this.

for i in range(len(alist)):
    for j in range(len(alist[i])):
        if alist[i][j].isWhatever:
            blist[i][j].doSomething()

Is there a nicer way to do something like this?

Asked By: Eugene M

||

Answers:

for d1 in alist
   for d2 in d1
      if d2 = "whatever"
          do_my_thing()
Answered By: Haoest

You could zip them. ie:

for a_row,b_row in zip(alist, blist):
    for a_item, b_item in zip(a_row,b_row):
        if a_item.isWhatever:
            b_item.doSomething()

However the overhead of zipping and iterating over the items may be higher than your original method if you rarely actually use the b_item (ie a_item.isWhatever is usually False). You could use itertools.izip instead of zip to reduce the memory impact of this, but its still probably going to be slightly slower unless you always need the b_item.

Alternatively, consider using a 3D list instead, so terrain for cell i,j is at l[i][j][0], objects at l[i][j][1] etc, or even combine the objects so you can do a[i][j].terrain, a[i][j].object etc.

[Edit] DzinX’s timings actually show that the impact of the extra check for b_item isn’t really significant, next to the performance penalty of re-looking up by index, so the above (using izip) seems to be fastest.

I’ve now given a quick test for the 3d approach as well, and it seems faster still, so if you can store your data in that form, it could be both simpler and faster to access. Here’s an example of using it:

# Initialise 3d list:
alist = [ [[A(a_args), B(b_args)] for i in xrange(WIDTH)] for j in xrange(HEIGHT)]

# Process it:
for row in xlist:
    for a,b in row:
        if a.isWhatever(): 
            b.doSomething()

Here are my timings for 10 loops using a 1000×1000 array, with various proportions of isWhatever being true are:

            ( Chance isWhatever is True )
Method      100%     50%      10%      1%

3d          3.422    2.151    1.067    0.824
izip        3.647    2.383    1.282    0.985
original    5.422    3.426    1.891    1.534
Answered By: Brian

Are you sure that the objects in the two matrices you are iterating in parallel are instances of conceptually distinct classes? What about merging the two classes ending up with a matrix of objects that contain both isWhatever() and doSomething()?

Answered By: Federico A. Ramponi

I’d start by writing a generator method:

def grid_objects(alist, blist):
    for i in range(len(alist)):
        for j in range(len(alist[i])):
            yield(alist[i][j], blist[i][j])

Then whenever you need to iterate over the lists your code looks like this:

for (a, b) in grid_objects(alist, blist):
    if a.is_whatever():
        b.do_something()
Answered By: Robert Rossney

Generator expressions and izip from itertools module will do very nicely here:

from itertools import izip
for a, b in (pair for (aline, bline) in izip(alist, blist) 
             for pair in izip(aline, bline)):
    if a.isWhatever:
        b.doSomething()

The line in for statement above means:

  • take each line from combined grids alist and blist and make a tuple from them (aline, bline)
  • now combine these lists with izip again and take each element from them (pair).

This method has two advantages:

  • there are no indices used anywhere
  • you don’t have to create lists with zip and use more efficient generators with izip instead.
Answered By: Dzinx

As a slight style change, you could use enumerate:

for i, arow in enumerate(alist):
    for j, aval in enumerate(arow):
        if aval.isWhatever():
            blist[i][j].doSomething()

I don’t think you’ll get anything significantly simpler unless you rearrange your data structures as Federico suggests. So that you could turn the last line into something like “aval.b.doSomething()”.

Answered By: John Fouhy

If the two 2D-lists remain constant during the lifetime of your game and you can’t enjoy Python’s multiple inheritance to join the alist[i][j] and blist[i][j] object classes (as others have suggested), you could add a pointer to the corresponding b item in each a item after the lists are created, like this:

for a_row, b_row  in itertools.izip(alist, blist):
    for a_item, b_item in itertools.izip(a_row, b_row):
        a_item.b_item= b_item

Various optimisations can apply here, like your classes having __slots__ defined, or the initialization code above could be merged with your own initialization code e.t.c. After that, your loop will become:

for a_row in alist:
    for a_item in a_row:
        if a_item.isWhatever():
            a_item.b_item.doSomething()

That should be more efficient.

Answered By: tzot

If anyone is interested in performance of the above solutions, here they are for 4000×4000 grids, from fastest to slowest:

EDIT: Added Brian’s scores with izip modification and it won by a large amount!

John’s solution is also very fast, although it uses indices (I was really surprised to see this!), whereas Robert’s and Brian’s (with zip) are slower than the question creator’s initial solution.

So let’s present Brian‘s winning function, as it is not shown in proper form anywhere in this thread:

from itertools import izip
for a_row,b_row in izip(alist, blist):
    for a_item, b_item in izip(a_row,b_row):
        if a_item.isWhatever:
            b_item.doSomething()
Answered By: Dzinx

When you are operating with grids of numbers and want really good performance, you should consider using Numpy. It’s surprisingly easy to use and lets you think in terms of operations with grids instead of loops over grids. The performance comes from the fact that the operations are then run over whole grids with optimised SSE code.

For example here is some numpy using code that I wrote that does brute force numerical simulation of charged particles connected by springs. This code calculates a timestep for a 3d system with 100 nodes and 99 edges in 31ms. That is over 10x faster than the best pure python code I could come up with.

from numpy import array, sqrt, float32, newaxis
def evolve(points, velocities, edges, timestep=0.01, charge=0.1, mass=1., edgelen=0.5, dampen=0.95):
    """Evolve a n body system of electrostatically repulsive nodes connected by
       springs by one timestep."""
    velocities *= dampen

    # calculate matrix of distance vectors between all points and their lengths squared
    dists = array([[p2 - p1 for p2 in points] for p1 in points])
    l_2 = (dists*dists).sum(axis=2)
    
    # make the diagonal 1's to avoid division by zero
    for i in xrange(points.shape[0]):
        l_2[i,i] = 1

    l_2_inv = 1/l_2
    l_3_inv = l_2_inv*sqrt(l_2_inv)

    # repulsive force: distance vectors divided by length cubed, summed and multiplied by scale
    scale = timestep*charge*charge/mass
    velocities -= scale*(l_3_inv[:,:,newaxis].repeat(points.shape[1], axis=2)*dists).sum(axis=1)

    # calculate spring contributions for each point
    for idx, (point, outedges) in enumerate(izip(points, edges)):
        edgevecs = point - points.take(outedges, axis=0)
        edgevec_lens = sqrt((edgevecs*edgevecs).sum(axis=1))
        scale = timestep/mass
        velocities[idx] += (edgevecs*((((edgelen*scale)/edgevec_lens - scale))[:,newaxis].repeat(points.shape[1],axis=1))).sum(axis=0)

    # move points to new positions
    points += velocities*timestep
Answered By: Ants Aasma

If a.isWhatever is rarely true you could build an “index” once:

a_index = set((i,j) 
              for i,arow in enumerate(a) 
              for j,a in enumerate(arow) 
              if a.IsWhatever())

and each time you want something to be done:

for (i,j) in a_index:
    b[i][j].doSomething()

If a changes over time, then you will need to
keep the index up-to-date. That’s why I used
a set, so items can be added and removed fast.

Answered By: user27030

A common pattern in other answers here is that they attempt to zip together the two inputs, then zip over elements from each pair of nested "row" lists. I propose to invert this, to get more elegant code. As the Zen of Python tells us, "Flat is better than nested."

I took the following approach to set up a test:

>>> class A:
...     def __init__(self):
...         self.isWhatever = True
... 
>>> 
>>> class B:
...     def doSomething(self):
...         pass
... 
>>> alist = [[A() for _ in range(1000)] for _ in range(1000)]
>>> blist = [[B() for _ in range(1000)] for _ in range(1000)]

Adapting the originally-best-performing code for 3.x, that solution was

def brian_modern():
    for a_row, b_row in zip(alist, blist):
        for a_item, b_item in zip(a_row, b_row):
            if a_item.isWhatever:
                b_item.doSomething()

(since nowadays, zip returns an iterator, and does what itertools.izip used to do).

On my platform (Python 3.8.10 on Linux Mint 20.3; Intel(R) Core(TM) i5-4430 CPU @ 3.00GHz with 8GB of DDR3 RAM @ 1600MT/s), I get this timing result:

>>> import timeit
>>> timeit.timeit(brian_modern, number=100)
10.740317705087364

Rather than this repeated zip, my approach is to flatten each input iterable first, and then zip the results.

from itertools import chain
def karl():
    flatten = chain.from_iterable
    for a_item, b_item in zip(flatten(alist), flatten(blist)):
        if a_item.isWhatever:
            b_item.doSomething()

This gives almost as good performance:

>>> karl()
>>> timeit.timeit(karl, number=100)
11.126002880046144

As a baseline, let’s try to pare the looping overhead to a minimum:

my_a = A()
my_b = B()
def baseline():
    a = my_a # avoid repeated global lookup
    b = my_b # avoid repeated global lookup
    for i in range(1000000):
        if a.isWhatever:
            b.doSomething()

and then check how much of the time is used by the actual object-checking logic:

>>> timeit.timeit(baseline, number=100)
9.41121925599873 

So, the pre-flattening approach does incur significantly more overhead (about 18%, vs. about 14% for the repeated-zip approach). However, it is still a fairly small amount of overhead, even for a trivial loop body, and it also allows us to write the code more elegantly.


In my testing, this is the fastest approach to pre-flattening. Splatting out arguments to itertools.chain is slightly slower again, while using a generator expression to flatten the input…

def karl_gen():
    a_flat = (i for row in alist for i in row)
    b_flat = (j for row in blist for j in row)
    for a_item, b_item in zip(a_flat, b_flat):
        if a_item.isWhatever:
            b_item.doSomething()

… is much slower:

>>> timeit.timeit(karl_gen, number=100)
16.904560427879915

Switching to list comprehensions here barely makes a difference to speed vs. the generators, while also doubling up on memory requirements temporarily. So itertools.chain.from_iterable is a clear winner.

Answered By: Karl Knechtel
Categories: questions Tags:
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.