Numpy: Сортировка многомерного массива многомерным массивом

Простите меня, если это избыточно или супер базово. Я прихожу на Python/Numpy из R и с трудом переворачиваю вещи в голове.

У меня есть n-мерный массив, который я хочу отсортировать, используя другой n-мерный массив значений индекса. Я знаю, что могу обернуть это в цикле, но кажется, что должен быть очень сжатый способ Numpyonic избить это в подчинении. Здесь мой примерный код для установки проблемы, где n = 2:

a1 = random.standard_normal(size=[2,5]) 
index = array([[0,1,2,4,3] , [0,1,2,3,4] ])

так что теперь у меня есть массив из 2 x 5 случайных чисел и индекс 2 x 5. Я прочитал справку для take() около 10 раз, но мой мозг, очевидно, не сдерживает ее.

Я подумал, что это может привести меня туда:

take(a1, index)

array([[ 0.29589188, -0.71279375, -0.18154864, -1.12184984,  0.25698875],
       [ 0.29589188, -0.71279375, -0.18154864,  0.25698875, -1.12184984]])

но это явно переупорядочивает только первый элемент (я полагаю, из-за уплощения).

Любые подсказки о том, как я получаю, откуда я пришел к решению, которое сортирует элемент 0 из a1 по элементу 0 индекса... element n?

Ответ 1

~~Я не могу придумать, как работать с этим в N измерениях, но~~ вот 2D-версия:

>>> a = np.random.standard_normal(size=(2,5))
>>> a
array([[ 0.72322499, -0.05376714, -0.28316358,  1.43025844, -0.90814293],
       [ 0.7459107 ,  0.43020728,  0.05411805, -0.32813465,  2.38829386]])
>>> i = np.array([[0,1,2,4,3],[0,1,2,3,4]]) 
>>> a[np.arange(a.shape[0])[:,np.newaxis],i]
array([[ 0.72322499, -0.05376714, -0.28316358, -0.90814293,  1.43025844],
       [ 0.7459107 ,  0.43020728,  0.05411805, -0.32813465,  2.38829386]])

Вот N-мерная версия:

>>> a[list(np.ogrid[[slice(x) for x in a.shape]][:-1])+[i]]

Вот как это работает:

Хорошо, начнем с 3-мерного массива для иллюстрации.

>>> import numpy as np
>>> a = np.arange(24).reshape((2,3,4))
>>> a
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]]])

Вы можете получить доступ к элементам этого массива, указав индекс вдоль каждой оси следующим образом:

>>> a[0,1,2]
6

Это эквивалентно a[0][1][2], так как вы будете обращаться к одному и тому же элементу, если бы мы имели дело со списком вместо массива.

Numpy позволяет вам стать еще более привлекательным при разрезании массивов:

>>> a[[0,1],[1,1],[2,2]]
array([ 6, 18])
>>> a[[0,1],[1,2],[2,2]]
array([ 6, 22])

Эти примеры были бы эквивалентны [a[0][1][2],a[1][1][2]] и [a[0][1][2],a[1][2][2]], если бы мы имели дело со списками.

Вы можете даже оставить повторяющиеся индексы, и numpy будет определять, что вы хотите. Например, приведенные выше примеры могут быть эквивалентно записаны:

>>> a[[0,1],1,2]
array([ 6, 18])
>>> a[[0,1],[1,2],2]
array([ 6, 22])

Форма массива (или списка), который вы срезаете в каждом измерении, влияет только на форму возвращаемого массива. Другими словами, numpy не заботится о том, чтобы вы пытались индексировать массив массивом формы (2,3,4), когда он вытягивает значения, за исключением того, что он вернет вам массив формы (2,3,4). Например:

>>> a[[[0,0],[0,0]],[[0,0],[0,0]],[[0,0],[0,0]]]
array([[0, 0],
       [0, 0]])

В этом случае мы снова и снова захватываем один и тот же элемент a[0,0,0], но numpy возвращает массив с той же формой, что и мы.

Хорошо, на вашу проблему. Вы хотите индексировать массив вдоль последней оси с числами в вашем массиве index. Итак, для примера в вашем вопросе вы хотели бы [[a[0,0],a[0,1],a[0,2],a[0,4],a[0,3]],a[1,0],a[1,1],...

Тот факт, что ваш индексный массив является многомерным, как я уже говорил ранее, не говорит о том, что вы хотите вывести эти индексы с нуля. он просто определяет форму выходного массива. Итак, в вашем примере вам нужно указать numpy, что первые 5 значений нужно вытащить из a[0], а последние 5 из a[1]. Легко!

>>> a[[[0]*5,[1]*5],index]

Он усложняется в N измерениях, но пусть это сделает для 3-мерного массива a, который я определил выше. Предположим, что мы имеем следующий индексный массив:

>>> i = np.array(range(4)[::-1]*6).reshape(a.shape)
>>> i
array([[[3, 2, 1, 0],
        [3, 2, 1, 0],
        [3, 2, 1, 0]],

       [[3, 2, 1, 0],
        [3, 2, 1, 0],
        [3, 2, 1, 0]]])

Итак, эти значения - все для индексов вдоль последней оси. Нам нужно указать numpy, какие индексы вдоль первой и второй осей должны быть взяты из этих чисел; то есть нам нужно указать numpy, что индексы для первой оси:

i1 = [[[0, 0, 0, 0],
       [0, 0, 0, 0],
       [0, 0, 0, 0]],

      [[1, 1, 1, 1],
       [1, 1, 1, 1],
       [1, 1, 1, 1]]]

и индексы для второй оси:

i2 = [[[0, 0, 0, 0],
       [1, 1, 1, 1],
       [2, 2, 2, 2]],

      [[0, 0, 0, 0],
       [1, 1, 1, 1],
       [2, 2, 2, 2]]]

Тогда мы можем просто сделать:

>>> a[i1,i2,i]
array([[[ 3,  2,  1,  0],
        [ 7,  6,  5,  4],
        [11, 10,  9,  8]],

       [[15, 14, 13, 12],
        [19, 18, 17, 16],
        [23, 22, 21, 20]]])

Удобная функция numpy, которая генерирует i1 и i2, называется np.mgrid. Я использую np.ogrid в своем ответе, который эквивалентен в этом случае из-за магии numpy, о которой я говорил ранее.

Надеюсь, что это поможет!

Ответ 2

После игры с этим еще сегодня я понял, что если бы я использовал функцию mapper вместе с take, я мог бы решить 2-мерную версию действительно так:

a1 = random.standard_normal(size=[2,5]) 
index = array([[0,1,2,4,3] , [0,1,2,3,4] ]) 
map(take, a1, index)

Мне нужно map() take() каждому элементу в a1

Конечно, принятый ответ решает n-мерную версию. Однако в ретроспективе я решил, что мне действительно не нужно n-мерное решение, а только двухмерная версия.