CODEWITHSUNDEEP
pythonarraysnumpyarray-broadcasting
konstant
konstant

Reputation: 705

How do I carry out this numpy array manipulation efficiently?

I am given with an array such as:

a = numpy.array([1, 2, 3, 4],[1, 2, 3, 4],[1, 2, 3, 4])

I am also given another array with all the entries greater than or equal to zero and less than the size of a. An example:

array_index = numpy.array([[0, 0, 1, 2], [0, 1, 2, 2], [0, 1, 1, 3]])

array_index basically tells where I should stack the elements of a to. For example first entry of array_index [0, 0, 1, 2] tells that index 0 element of fist entry of a should stay the same, second index 0 tells that index 1 element of a should be stacked to index 0, third index 1 tells that index 2 element of a should be stacked to index 1, and so on. Now I need to construct another array of the same shape as that of my_array such that the entries are the respective index value of a at array_index. In this case I would have first

array_desired = numpy.zeros(array_index.shape)

And I would fill in the values like:

array_desired = [[a[0]+a[1], a[2], a[3], 0], [a[0], a[1], a[2]+a[3], 0],
                 [a[0], a[1]+a[2], 0, a[3]]]

To do this, as a first approach I constructed the following for the 0th element:

stack_index = numpy.where(array_index == 0)
array_desired[stack_index] += a[stack_index]
print(array_desired)
>>> [[1 2 0 0] [1 0 0 0] [1 0 0 0]]

which is something, but not exactly what I wanted:

>>>[[3 0 0 0][1 0 0 0][1 0 0 0]]

Any idea on how to achieve proper stacking by index as I described above?

Update:

I have a method that works for first index now:

temp_array = numpy.zeros(array_index.shape)
stack_index = numpy.where(array_index == 0)
temp_array[stack_index] = a[stack_index]
n = np.sum(temp_array, axis=(1))
array_desired[:,0]= n

But, this method would still require me to loop over each index. I would like to have something more efficient.

Upvotes: 1

Views: 96

Answers (1)

Warren Weckesser
Warren Weckesser

Reputation: 114911

I think this does what you want.

Here's your sample data:

In [98]: a
Out[98]: 
array([[1, 2, 3, 4],
       [1, 2, 3, 4],
       [1, 2, 3, 4]])

In [99]: array_index
Out[99]: 
array([[0, 0, 1, 2],
       [0, 1, 2, 2],
       [0, 1, 1, 3]])

Create b to hold the result, and do the sum using np.add.at. row_index is just the 3x1 array [[0], [1], [2]].

In [100]: b = np.zeros_like(a)

In [101]: row_index = np.arange(array_index.shape[0]).reshape(-1, 1)

In [102]: np.add.at(b, (row_index, array_index), a)

In [103]: b
Out[103]: 
array([[3, 3, 4, 0],
       [1, 2, 7, 0],
       [1, 5, 0, 4]])

Upvotes: 2

Related Questions