Add repeated elements of array indexed by another array

Question

I have a relatively simple problem that I cannot solve without using loops. It is difficult for me to figure out the correct title for this problem. Lets say we have two numpy arrays:

array_1 = np.array([[0, 1, 2],
                    [3, 3, 3],
                    [3, 3, 4],
                    [3, 6, 2]])

array_2 = np.array([[0, 0, 0], 
                    [1, 1, 1],
                    [2, 2, 2],
                    [3, 3, 3],
                    [4, 4, 4],
                    [5, 5, 5],
                    [6, 6, 6]])

array_1 represents indices of the rows in array_2 that we want to sum. So for example, 4th row in result array should contain summed all rows in array_2 that have same row indices as all 3s in array_1.

It is much easier to understand it in the code:

result = np.empty(array_2.shape)

for i in range(array_1.shape[0]):
    for j in range(array_1.shape[1]):
        index = array_1[i, j]
        result[index] = result[index] + array_2[i]

Result should be:

[[ 0  0  0]
 [ 0  0  0]
 [ 3  3  3]
 [10 10 10]
 [ 2  2  2]
 [ 0  0  0]
 [ 3  3  3]]

I tried to use np.einsum but I need to use both elements in array as indices and also its rows as indices so I'm not sure if np.einsum is the best path here.

This is the problem I have in graphics. array_1 represent indices of vertices for triangles and array_2 represents normals where index of a row corresponds to the index of the vertex

Problem is clear, but I am curious: why would you want to solve this without loops? — sarema
– sarema, Commented Jan 11, 2022 at 21:16
This is the problem I have in graphics so array_1 and array_2 can be massive and this can get extremely slow. But what bugs me the most with this problem is that there are just two simple for loops and I can't port them to numpy. It is making my brain boil for 3rd day — Nemanja Stojanovic
– Nemanja Stojanovic, Commented Jan 11, 2022 at 21:19
"and I can't port them to numpy.". Why do you want to do this? Do you want to speed up the algorithm? What is it you are trying to achieve at the end of the day? — sarema
– sarema, Commented Jan 11, 2022 at 21:21
Get the looped version working correctly and show the results. Then we can talk about doing the same thing with the full power of numpy. — hpaulj
– hpaulj, Commented Jan 11, 2022 at 21:53

Mad Physicist · Accepted Answer · 2022-01-12 03:51:20Z

5

Any time you're adding something from a repeated index, normal ufuncs like np.add don't work out of the box because they only process a repeated fancy index once. Instead, you have to use the unbuffered version, which is np.add.at.

Here, you have a pair of indices: the row in array_1 is the row index into array_2, and the element of array_1 is the row index into the output.

First, construct the indices explicitly as fancy indices. This will make it much simpler to use them:

output_row = array_1.ravel()
input_row = np.repeat(np.arange(array_1.shape[0]), array_1.shape[1]).ravel()

You can apply input_row directly to array_2, but you need add.at to use output_row:

output = np.zeros_like(array_2)
np.add.at(output, output_row, array_2[input_row])

You really only use the first four rows of array_2, so it could be truncated to

array_2 = array2[:array_1.shape[0]]

In that case, you would want to initialize the output as:

output = np.zeros_like(array_2, shape=(output_row.max() + 1, array2.shape[1]))

edited Jan 12, 2022 at 3:51

answered Jan 11, 2022 at 21:32

Mad Physicist

116k29 gold badges202 silver badges292 bronze badges

Sign up to request clarification or add additional context in comments.

5 Comments

sarema Over a year ago

That's not correct. The shape of the result is (4, 3), which is not what the OP is looking for.

Mad Physicist Over a year ago

@sarema. See update. I have the same answer as OP now.

Nemanja Stojanovic Over a year ago

Wow. I will have to reread this answer a couple of times to fully understand it. Thank You!. Anyway, I have no idea what title to put on this question. Any help would be appreciated!

sarema Over a year ago

I see. Impressive!

Mad Physicist Over a year ago

@NemanjaStojanovic. I've found fancy indexing to be much more counterintuitive than np.ufunc.at. All it does is allow you to do output[output_row] += array_2[input_row] correctly. I would recommend two things: first forget that the arrays are 2D and work with a 1D example, then read the docs on fancy indexing: numpy.org/doc/stable/user/…

Collectives™ on Stack Overflow

Add repeated elements of array indexed by another array

1 Answer 1

5 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

5 Comments

Your Answer

Sign up or log in

Post as a guest

Related