[SciPy-User] "inverting" an array

[nicky van foreest]

Hi,

I am wondering whether a shortcut exists in numpy/scipy for the following
problem.  The values in an array represent the number of customers that
arrive in a certain time slot, e.g.,

a = [0,4,7,3,1,5, 0,0,0,]

means that in time slot 1 4 customers arrive, in time slot 2 seven arrive,
and so on. Now I want to "invert" this array to compute the arrival time of
the i-th customer. Thus, customer 2 arrives in time slot 1, customer 6 in
time slot 2, and so on. For this problem I wrote the following function:

a = [0,4,7,3,1,5, 0,0,0,]
A = np.cumsum(a)

def invert(A):
    Ainv = np.empty(A[-1])
    aprev=0
    for i, a in enumerate(A):
        Ainv[aprev:a] = i
        aprev = a
    return Ainv


Ainv= invert(A)

print a
print A
print Ainv

The output is

[0, 4, 7, 3, 1, 5, 0, 0, 0]
[ 0  4 11 14 15 20 20 20 20]
[ 1.  1.  1.  1.  2.  2.  2.  2.  2.  2.  2.  3.  3.  3.  4.  5.  5.  5.
  5.  5.]

Does anybody know whether this code can be made faster, or whether a
numpy/scipy function exists that establishes this in one go?

thanks

Nicky
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mail.scipy.org/pipermail/scipy-user/attachments/20140204/f6ea1383/attachment.html>