Numpy array index handling

[phishboh at gmail.com]

Being a Matlab user wanting to switch to Python/SciPy, I'd like to
know how the following Matlab code would be written in Python:

% First, just some artificial data
N = 1000000 ;
input = sign(randn(1, N)) ;
a = (1 : N) ;
% This is what I'd like to do;
% for indices where the input has a certain value, compute and store
something
% for indices where the input has another certain value, compute and
store something else
output(input < 0) = 10 * a(input < 0) ;
output(input > 0) = 20 * a(input > 0) ;

I have tried the following in Python:
N = 1000000
input = sign(randn(N))
a = arange(N)
output = zeros(N)
output[input < 0] = 10 * a[input < 0]
output[input > 0] = 20 * a[input > 0]

However, that gives me in IndexError.

Of course I could use a for-loop, but I assume that there is another
way (although I haven't been able to find it in any documentation).

Many thanks in advance!