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

Martin Spacek numpy@mspacek.mm...
Fri Aug 6 16:39:41 CDT 2010

On 2010-08-06 13:11, Martin Spacek wrote:
> Josef, I'd forgotten you could use None to increase the dimensionality of an
> array. Neat. And, somehow, it's almost twice as fast as the Cython version!:
>
>   >>>  timeit a[np.arange(a.shape[0])[:, None], i]
> 100000 loops, best of 3: 5.76 us per loop

I just realized why the Cython version was slower - the two assertion lines.
commenting those out:

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True) # might be necessary to release the GIL?
def rowtake_cy(np.ndarray[np.int32_t, ndim=2] a,
np.ndarray[np.int32_t, ndim=2] i):
"""For each row in a, return values according to column indices in the
corresponding row in i. Returned shape == i.shape"""

cdef Py_ssize_t nrows, ncols, rowi, coli
cdef np.ndarray[np.int32_t, ndim=2] out

nrows = i.shape[0]
ncols = i.shape[1] # num cols to take for each row
#assert a.shape[0] == nrows
#assert i.max() < a.shape[1]
out = np.empty((nrows, ncols), dtype=np.int32)

for rowi in range(nrows):
for coli in range(ncols):
out[rowi, coli] = a[rowi, i[rowi, coli]]

return out

gives me:

>>> timeit rowtake_cy(a, i)
1000000 loops, best of 3: 1.44 us per loop

which is 4X faster than the a[np.arange(a.shape[0])[:, None], i] broadcasting
method.

Martin