Most Pythonic way of assigning values of array to another array

user_4265321 0
Viewed 54

I have two arrays with the same length (6, in this example). One stores floats:

a = np.array([0.2, 0.01, 0.5, 0.7, 0., 0.002])

the second one stores indices (hence, int values):

indices = np.array([4, 9, 0, 2, 2, 4])

In my code I initialize another array, which has length in general different from that of a and indices, for example 10 in this example:

c = np.zeros(10)

I would like to find a Pythonic way to accomplish the following:

for i in range(len(indices)):
    c[indices[i]] += a[i]

which, in this example, produces:

[0.5   0.    0.7   0.    0.202 0.    0.    0.    0.    0.01 ]

I tried looking at this brilliant example, however I am not sure how this can be applied here.

unique for-loop numpy python
2 Answers

You can use the .at method of the np.add ufunc:

np.add.at(c, indices, a)

Here's the help page for the .at method of ufuncs:

at(...) method of numpy.ufunc instance
    at(a, indices, b=None, /)

    Performs unbuffered in place operation on operand 'a' for elements
    specified by 'indices'. For addition ufunc, this method is equivalent to
    ``a[indices] += b``, except that results are accumulated for elements that
    are indexed more than once. For example, ``a[[0,0]] += 1`` will only
    increment the first element once because of buffering, whereas
    ``add.at(a, [0,0], 1)`` will increment the first element twice.

    .. versionadded:: 1.8.0

    Parameters
    ----------
    a : array_like
        The array to perform in place operation on.
    indices : array_like or tuple
        Array like index object or slice object for indexing into first
        operand. If first operand has multiple dimensions, indices can be a
        tuple of array like index objects or slice objects.
    b : array_like
        Second operand for ufuncs requiring two operands. Operand must be
        broadcastable over first operand after indexing or slicing.
user_50144550

For sum, your operation is exactly what bincount is:

np.bincount(indices, weights=a)

Output:

array([0.5  , 0.   , 0.7  , 0.   , 0.202, 0.   , 0.   , 0.   , 0.   ,  0.01 ])
user_42384080
Related