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 <[email protected]> wrote: > On Thu, Aug 5, 2010 at 1:32 PM, <[email protected]> wrote: > > On Thu, Aug 5, 2010 at 4:07 PM, Martin Spacek <[email protected]> > 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 > [email protected] > http://mail.scipy.org/mailman/listinfo/numpy-discussion >
_______________________________________________ NumPy-Discussion mailing list [email protected] http://mail.scipy.org/mailman/listinfo/numpy-discussion
