Return generator instead of list of locations from a 2d Array

Question

I was working on a game yesterday where I had to traverse a 2-D array and find the location of any cells marked "d" (where cells are represented as either "-" for blank or "d" for dirty).

I did this using two for-loops:

def find_dirty_cells(board):

  dirty_cells = []
    for enum, row in enumerate(board):
      for enumrow, cell in enumerate(row):
        if cell == 'd':
          dirty_cells.append((enum, enumrow))

  return dirty_cells 

But then I thought it might be better to build a generator object and return that instead, so I wrote this:

def find_dirty_cells(board):
  return ((enum, enumrow) for enumrow, cell in enumerate(row) for enum, row in enumerate(board) if cell == 'd')

But the second gives me incorrect data in response (i.e., it doesn't find the 'd' cells). There must be something simple I am overlooking that makes the second not equal to the first, but I can't see it. The real question I tried to solve is this: is there an easy way to make my first attempt return a generator?

Answer 1

You need to list your for loops in the same order you nest them:

def find_dirty_cells(board):
    return ((enum, enumrow) for enum, row in enumerate(board) for enumrow, cell in enumerate(row) if cell == 'd')

You had swapped the enumerate(board) and enumerate(row) loops, which would only work if there was a global row variable still available in your session or in your module.

It may be easier in this case to just use a generator function instead of a generator expression, by using yield instead of dirty_cells.append():

def find_dirty_cells(board):
    for enum, row in enumerate(board):
        for enumrow, cell in enumerate(row):
            if cell == 'd':
                yield enum, enumrow

This will have the exact same effect but is perhaps more readable.

Answer 2

To cleanly convert your original function into a generator, what you want is the yield statement, rather than the return (or in your particular case, rather than the append). I'd prefer this version to the generator expression version, because the original is vastly more readable.

def find_dirty_cells(board):
    for enum, row in enumerate(board):
        for enumrow, cell in enumerate(row):
            if cell == 'd':
                yield (enum, enumrow)