[Numpy-discussion] use index array of len n to select columns of n x m array

John Salvatier jsalvati@u.washington....
Thu Aug 5 16:00:19 CDT 2010


choose might be slower if you weren't doing an "arange(N)" each time.

On Thu, Aug 5, 2010 at 1:51 PM, Keith Goodman <kwgoodman@gmail.com> wrote:

> On Thu, Aug 5, 2010 at 1:32 PM,  <josef.pktd@gmail.com> wrote:
> > On Thu, Aug 5, 2010 at 4:07 PM, Martin Spacek <numpy@mspacek.mm.st>
> wrote:
> >> josef.pkt wrote:
> >>>>> a = np.array([[0, 1],
> >>                   [2, 3],
> >>                   [4, 5],
> >>                   [6, 7],
> >>                   [8, 9]])
> >>>>> i = np.array([0, 1, 1, 0, 1])
> >>>>> a[range(a.shape[0]), i]
> >> array([0, 3, 5, 6, 9])
> >>>>> a[np.arange(a.shape[0]), i]
> >> array([0, 3, 5, 6, 9])
> >>
> >>
> >> Thanks for all the tips. I guess I was hoping for something that could
> avoid
> >> having to generate np.arange(a.shape[0]), but
> >>
> >>  >>> a[np.arange(a.shape[0]), i]
> >>
> >> sure is easy to understand. Is there maybe a more CPU and/or memory
> efficient
> >> way? I kind of like John Salvatier's idea:
> >>
> >>  >>> np.choose(i, (a[:,0], a[:,1])
> >>
> >> but that would need to be generalized to "a" of arbitrary columns.
> >
> > seems to work:
> >
> >>>> np.choose(i, a.T)
> > array([0, 3, 5, 6, 9])
> >
> > but wouldn't get around the 31 limit that you found.
>
> Choose is fast:
>
> >> N = 1000
> >> a = np.random.randint(0, 9, (N,2))
> >> i = np.random.randint(0, 2, N)
> >> timeit a[range(N), i]
> 10000 loops, best of 3: 108 us per loop
> >> timeit a[np.arange(N), i]
> 10000 loops, best of 3: 39 us per loop
> >> timeit np.choose(i, a.T)
> 10000 loops, best of 3: 32.3 us per loop
>
> But flat is faster:
>
> >> timeit idx = np.arange(N); idx *= 2; idx += i; a.flat[idx]
> 100000 loops, best of 3: 16.8 us per loop
> _______________________________________________
> NumPy-Discussion mailing list
> NumPy-Discussion@scipy.org
> http://mail.scipy.org/mailman/listinfo/numpy-discussion
>
-------------- next part --------------
An HTML attachment was scrubbed...
URL: http://mail.scipy.org/pipermail/numpy-discussion/attachments/20100805/bad114ab/attachment-0001.html 


More information about the NumPy-Discussion mailing list