Broadcast numpy array values

user_14918720 0
Viewed 73

I have a primary two-dimensional numpy array with multiple duplicates and a secondary one with the unique values of the first.

[[ 0  0  1 ]
 [ 1  0  2 ]
 [ 2  0  2 ]
 ...
 [ 0  0  1 ]
 [ 3  0  2 ]
 [ 2  0  2 ]]

[[ 0  0  1 ]
 [ 1  0  2 ]
 [ 2  0  2 ]
 [ 3  0  2 ]]

A third two-dimensional array (with the same number of rows as the second) contains some extra values, which are produced from the corresponding row in the second array.

[[ 2 -0.5]
 [ 2.4  0.5]
 [ 2.5  0.4]
 [ 2.4  1]]

What I have in mind is to create new array using the values of the first as keys and the values of the third as the content. The order has to be maintained.

[[ 2 -0.5]
 [ 2.4  0.5]
 [ 2.5  0.4]
 ...
 [ 2 -0.5]
 [ 2.4  1]
 [ 2.5  0.4]]

I tried creating pandas dfs (one from the first array and one from the second+third) and merging them on the first 3 columns, but the result is way too slow for my needs.

pandas numpy python
1 Answers

In general you can solve this using broadcasted equality to find the appropriate indices, for example:

x = np.array([[ 0,  0,  1 ],
              [ 1,  0,  2 ],
              [ 2,  0,  2 ],
              [ 0,  0,  1 ],
              [ 3,  0,  2 ],
              [ 2,  0,  2 ]])

x_unique = np.array([[ 0,  0,  1 ],
                     [ 1,  0,  2 ],
                     [ 2,  0,  2 ],
                     [ 3,  0,  2 ]])

values = np.array([[ 2, -0.5],
                   [ 2.4,  0.5],
                   [ 2.5,  0.4],
                   [ 2.4,  1]])

index = np.argmax((x[:, None] == x_unique).all(-1), axis=1)
values[index]
# array([[ 2. , -0.5],
#        [ 2.4,  0.5],
#        [ 2.5,  0.4],
#        [ 2. , -0.5],
#        [ 2.4,  1. ],
#        [ 2.5,  0.4]])

If you are generating the unique values using np.unique, you can do this more efficiently by computing the indices when you extract the unique values:

x_unique2, index = np.unique(x, axis=0, return_inverse=True)
assert np.allclose(x_unique, x_unique2)  # same as above
print(values[index])
# array([[ 2. , -0.5],
#        [ 2.4,  0.5],
#        [ 2.5,  0.4],
#        [ 2. , -0.5],
#        [ 2.4,  1. ],
#        [ 2.5,  0.4]])
user_29378310
Related