Взять элементы массива на основе индексов в другом массиве - PullRequest
Новая
       44

Взять элементы массива на основе индексов в другом массиве

0 голосов
/ 13 января 2020

У меня есть массив значений на одной стороне: 

A = np.arange(30).reshape((3, 10))

Out: array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
            [10, 11, 12, 13, 14, 15, 16, 17, 18, 19],
            [20, 21, 22, 23, 24, 25, 26, 27, 28, 29]])

И массив индексов, ссылающихся на него, где каждый столбец ссылается на каждую строку в A. 

np.random.seed(0)
index = np.random.randint(0, 9, 6).reshape((2, 3))

Out: array([[5, 0, 3],
            [3, 7, 3]])

Я хочу получить массив с такими же размерами массива индекса, но заменив каждый индекс его значением в A. Я смог выполнить sh на: 

 np.dstack([A[0].take(index.T[0]),        
            A[1].take(index.T[1]),        
            A[2].take(index.T[2])]).squeeze()

Out: array([[ 5, 10, 23],
            [ 3, 17, 23]])

Мне кажется, что я что-то упустил, и это не оптимальный способ сделать это. Я также обеспокоен производительностью, когда размер массивов увеличивается. Есть ли более общий c и масштабируемый способ достижения sh этого?

1 Ответ

3 голосов
/ 13 января 2020

Вы можете использовать np.take_along_axis: 

np.take_along_axis(A, index.T, 1).T

array([[ 5, 10, 23],
       [ 3, 17, 23]])
...