Vectorizing a Numpy slice operation

Issue

Say I have a Numpy vector,

A = zeros(100)

and I divide it into subvectors by a list of breakpoints which index into A, for instance,

breaks = linspace(0, 100, 11, dtype=int)

So the i-th subvector would be lie between the indices breaks[i] (inclusive) and breaks[i+1] (exclusive).
The breaks are not necessarily equispaced, this is only an example.
However, they will always be strictly increasing.

Now I want to operate on these subvectors. For instance, if I want to set all elements of the i-th subvector to i, I might do:

for i in range(len(breaks) - 1):
    A[breaks[i] : breaks[i+1]] = i

Or I might want to compute the subvector means:

b = empty(len(breaks) - 1)
for i in range(len(breaks) - 1):
    b = A[breaks[i] : breaks[i+1]].mean()

And so on.

How can I avoid using for loops and instead vectorize these operations?

Solution

There really isn’t a single answer to your question, but several techniques that you can use as building blocks. Another one you may find helpful:

All numpy ufuncs have a .reduceat method, which you can use to your advantage for some of your calculations:

>>> a = np.arange(100)
>>> breaks = np.linspace(0, 100, 11, dtype=np.intp)
>>> counts = np.diff(breaks)
>>> counts
array([10, 10, 10, 10, 10, 10, 10, 10, 10, 10])
>>> sums = np.add.reduceat(a, breaks[:-1], dtype=np.float)
>>> sums
array([  45.,  145.,  245.,  345.,  445.,  545.,  645.,  745.,  845.,  945.])
>>> sums / counts  # i.e. the mean
array([  4.5,  14.5,  24.5,  34.5,  44.5,  54.5,  64.5,  74.5,  84.5,  94.5])

Answered By – Jaime

This Answer collected from stackoverflow, is licensed under cc by-sa 2.5 , cc by-sa 3.0 and cc by-sa 4.0

Leave a Reply

(*) Required, Your email will not be published