[Numpy-discussion] Can you enumerate an array?

[Chris Barker]

Arnd,

You've written some nice compact code, but it only works for rank-2 
arrays. Did you see Travis' solution? I think he's put it (or a C 
version, I'm not sure about that) in scipy_base. However, not using flat 
is nice, as in Numeric and numarray is sometimes makes a copy, and 
sometimes doesn't.

-Chris

Travis Oliphant wrote:
> class ndenumerate(enumerate):
>    def __init__(self, arr):
>        self.iter = enumerate(arr.flat)
>        self.ashape = arr.shape
>        self.nd = arr.ndim
>        self.tot = arr.size
>        self.factors = [None]*(self.nd-1)
>        val = self.ashape[-1]
>        for i in range(self.nd-1,0,-1):
>            self.factors[i-1] = val
>            val *= self.ashape[i-1]
> 
>    def __len__(self):
>        return self.tot
> 
>    def next(self):
>        res = self.iter.next()
>        indxs = [None]*self.nd
>        val = res[0]
>        for i in range(self.nd-1):
>            indxs[i] = val / self.factors[i]
>            val = val % self.factors[i]
>        indxs[self.nd-1] = val
>        return tuple(indxs), res[1]

-- 
Christopher Barker, Ph.D.
Oceanographer
                                     		
NOAA/OR&R/HAZMAT         (206) 526-6959   voice
7600 Sand Point Way NE   (206) 526-6329   fax
Seattle, WA  98115       (206) 526-6317   main reception

Chris.Barker at noaa.gov