Leon Sit
Leon Sit

Reputation:

dual iterator in one python object

In python, I am trying to write a class that support two different kind of iterator. Roughly speaking, this object contains a matrix of data and I want to have two different kind of iterator to support row iteration and column iteration.

Upvotes: 2

Views: 911

Answers (3)

Alex Martelli
Alex Martelli

Reputation: 882681

dict has several iterator-producing methods -- iterkeys, itervalues, iteritems -- and so should your class. If there's one "most natural" way of iterating, you should also alias it to __iter__ for convenience and readability (that's probably going to be iterrows; of course there is always going to be some doubt, as there was with dict when we designed its iteration behavior, but a reasonable pick is better than none).

For example, suppose your matrix is square, held flattened up into a row-major list self.data, with a side of self.n. Then:

def iterrows(self):
  start = 0
  n = self.n
  data = self.data
  while start < n*n:
    stop = start + n
    yield data[start:stop]
    start = stop

def itercols(self):
  start = 0
  n = self.n
  data = self.data
  while start < n:
    yield data[start::n]
    start += 1

__iter__ = iterrows

Upvotes: 5

Jon-Eric
Jon-Eric

Reputation: 17295

Is this what you're looking for?

class Matrix(object):
    def __init__(self, rows):
        self._rows = rows

    def columns(self):
        return zip(*self._rows)

    def rows(self):
        return self._rows

# Create a Matrix by providing rows.
m = Matrix([[1,2,3],
            [4,5,6],
            [7,8,9]])

# Iterate in row-major order.
for row in m.rows():
    for value in row:
        print value

# Iterate in column-major order.
for column in m.columns():
    for value in column:
        print value

You can use itertools.izip instead of zip if you want to create each column on demand.

You can also move the iteration of the actual values into the class. I wasn't sure if you wanted to iterate over rows/columns (as shown) or the values in the rows/columns.

Upvotes: 4

habnabit
habnabit

Reputation: 10294

Okay, so make two separate methods where each is a generator.

class Matrix(object):
    def iter_rows(self):
        for row in self.rows:
            yield row

    def iter_columns(self):
        for column in self.columns:
            yield column

Your __iter__ could iterate over one or the other by default, though I'd recommend having no __iter__ at all.

Upvotes: 2

Related Questions