Search code examples
pythonnumpyfor-loopunique

Most Pythonic way of assigning values of array to another array


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.


Solution

  • 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.