[Numpy-discussion] Why does argwhere return column vector?

Stéfan van der Walt stefan@sun.ac...
Tue May 27 04:43:47 CDT 2008


Hi Mark

2008/5/27 mark <markbak@gmail.com>:
> OK, I get how it works for 2D arrays.
>
> What I want to do is insert a number, say 7, before every value in the
> array that is larger than, for example, 1.
>
> Then I need to first find all the indices of values larger than 1, and
> then I can do an insert:
>
>>>> a = arange(5)
>>>> i = argwhere( a>1 )
>>>> insert(a,i[:,0],7)
> array([0, 1, 7, 2, 7, 3, 7, 4])
>
> Is there a better way to do this?

Inserting is slow, since a new array is allocated each time.  In the
case where you insert
at multiple indexes, it may be handled better, but I haven't timed it.

I would do it the following way:

import numpy as np
a = np.array([2,1,3,-1,2])

mask = a > 1
out = np.zeros(len(a) + sum(mask), dtype=int)
out.fill(7)
out[np.arange(len(a)) + mask.cumsum()] = a

Regards
Stéfan


More information about the Numpy-discussion mailing list