Numpy: объединить список массивов другим массивом (альтернатива np.choose) - PullRequest
1 голос
/ 23 марта 2019

У меня есть список массивов, каждый из которых имеет одинаковую форму.Скажем:

a = [np.array([[1, 2, 3],
               [4, 5, 6],
               [7, 8, 9]]),
     np.array([[11, 12, 13],
               [14, 15, 16],
               [17, 18, 19]]),
     np.array([[99, 98, 97],
               [96, 95, 94],
               [93, 92, 91]])]

И у меня есть еще один массив той же формы, который дает индексы списка, из которых я хочу получить элементы:

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

То, что я хочу получить, этоследующее:

np.array([[1, 2, 13],
          [96, 15, 6],
          [93, 18, 91]])

Было простое решение, которое работало нормально:

np.choose(b, a)

Но это ограничено максимум 32 массивами.Но в моем случае мне нужно объединить больше массивов (более 100).Поэтому мне нужен другой способ сделать это.

Я полагаю, это должно быть что-то, касающееся индексации достижений или, возможно, np.take метода.Поэтому, вероятно, первый шаг - a = np.array(a), а затем что-то вроде a[np.arange(a.shape[0]), b].Но у меня не получается.

Может кто-нибудь помочь?:)

Ответы [ 2 ]

2 голосов
/ 23 марта 2019

Вы можете попробовать использовать np.ogrid.На основании этого ответа.Конечно, вам придется сначала преобразовать a в массив NumPy

i, j = np.ogrid[0:3, 0:3]
print (a[b, i, j])

# array([[ 1,  2, 13],
#        [96, 15,  6],
#        [93, 18, 91]])
1 голос
/ 23 марта 2019
In [129]: a = [np.array([[1, 2, 3], 
     ...:                [4, 5, 6], 
     ...:                [7, 8, 9]]), 
     ...:      np.array([[11, 12, 13], 
     ...:                [14, 15, 16], 
     ...:                [17, 18, 19]]), 
     ...:      np.array([[99, 98, 97], 
     ...:                [96, 95, 94], 
     ...:                [93, 92, 91]])]                                        
In [130]: b = np.array([[0, 0, 1], 
     ...:               [2, 1, 0], 
     ...:               [2, 1, 2]])                                             
In [131]:                                                                       
In [131]: A = np.array(a)                                                       
In [132]: A.shape                                                               
Out[132]: (3, 3, 3)

Вы хотите использовать b для индексации первого измерения. Для других измерений вам нужны индексы, которые транслируются с b, то есть вектор столбца и вектор строки:

In [133]: A[b, np.arange(3)[:,None], np.arange(3)]                              
Out[133]: 
array([[ 1,  2, 13],
       [96, 15,  6],
       [93, 18, 91]])

Существуют различные вспомогательные функции для создания этих массивов, например,

In [134]: np.ix_(range(3),range(3))                                             
Out[134]: 
(array([[0],
        [1],
        [2]]), array([[0, 1, 2]]))

и ogrid, как указано в другом ответе.

Вот относительно новая функция, которая также выполняет свою работу:

In [138]: np.take_along_axis(A, b[None,:,:], axis=0)                              
Out[138]: 
array([[[ 1,  2, 13],
        [96, 15,  6],
        [93, 18, 91]]])

Мне пришлось немного подумать, прежде чем я получил правильную настройку b.

...