Issue
I have 2 arrays as input. On array as output. Array a
holds the data and is of shape (N,M)
, while array b
holds the indices and is of shape (N,X,2)
. The resulting array should be of shape (N,X)
, with the values taken from a
.
Right now it only works with a for loop. How could I vectorize it since I have huge arrays as input?
Below is a sample code to demonstrate what I have right now:
import numpy as np
# a of shape (N,M)
# b of shape (N,X,2)
# t_result of shape (N, X)
a = np.random.randint(0, 10, size=(10, 10))
b = np.random.randint(0, 2, size=(10, 9, 2))
t_result = np.empty((10, 9))
for i in range(b.shape[0]):
t_result[i] = a[i, b[i, :, 0]]
print(t_result)
print(t_result.shape)
Solution
Ok so I adapted a bit the answer to another post from scleronomic:
import numpy as np
# a of shape (N,M)
# b of shape (N,X,2)
# t_result of shape (N, X)
a = np.random.randint(0, 10, size=(10, 10))
b = np.random.randint(0, 2, size=(10, 9, 2))
t_result = np.empty((10, 9))
t_result = a[np.arange(a.shape[0])[:,None],b[np.arange(b.shape[0]),:,0]]
print(t_result)
print(t_result.shape)
I am not sure whether or not it is the best solution but it works.
Answered By - Gildur7161
0 comments:
Post a Comment
Note: Only a member of this blog may post a comment.