Как использовать numpy.argsort() в качестве индексов в более чем двух измерениях?

Я знаю, что что-то похожее на этот вопрос уже задавалось много раз, но все ответы на подобные вопросы, похоже, работают только для массивов с двумя измерениями.

Мое понимание np.argsort() в том, что np.sort(array) == array[np.argsort(array)] должно быть True, Я обнаружил, что это действительно правильно, если np.ndim(array) == 2, но это дает разные результаты, если np.ndim(array) > 2,

Пример:

>>> array = np.array([[[ 0.81774634,  0.62078744],
                       [ 0.43912609,  0.29718462]],
                      [[ 0.1266578 ,  0.82282054],
                       [ 0.98180375,  0.79134389]]])
>>> np.sort(array)
array([[[ 0.62078744,  0.81774634],
        [ 0.29718462,  0.43912609]],

       [[ 0.1266578 ,  0.82282054],
        [ 0.79134389,  0.98180375]]])
>>> array.argsort()
array([[[1, 0],
        [1, 0]],

       [[0, 1],
        [1, 0]]])
>>> array[array.argsort()]
array([[[[[ 0.1266578 ,  0.82282054],
          [ 0.98180375,  0.79134389]],

         [[ 0.81774634,  0.62078744],
          [ 0.43912609,  0.29718462]]],


        [[[ 0.1266578 ,  0.82282054],
          [ 0.98180375,  0.79134389]],

         [[ 0.81774634,  0.62078744],
          [ 0.43912609,  0.29718462]]]],



       [[[[ 0.81774634,  0.62078744],
          [ 0.43912609,  0.29718462]],

         [[ 0.1266578 ,  0.82282054],
          [ 0.98180375,  0.79134389]]],


        [[[ 0.1266578 ,  0.82282054],
          [ 0.98180375,  0.79134389]],

         [[ 0.81774634,  0.62078744],
          [ 0.43912609,  0.29718462]]]]])

Итак, кто-нибудь может объяснить мне, как именно np.argsort() можно использовать как индексы для получения отсортированного массива? Единственный способ, которым я могу придумать это:

args = np.argsort(array)
array_sort = np.zeros_like(array)
for i in range(array.shape[0]):
    for j in range(array.shape[1]):
        array_sort[i, j] = array[i, j, args[i, j]]

который чрезвычайно утомителен и не может быть обобщен для любого данного числа измерений.

Источник

3 ответа

Решение

Вот общий метод:

import numpy as np

array = np.array([[[ 0.81774634,  0.62078744],
                   [ 0.43912609,  0.29718462]],
                  [[ 0.1266578 ,  0.82282054],
                   [ 0.98180375,  0.79134389]]])

a = 1 # or 0 or 2

order = array.argsort(axis=a)

idx = np.ogrid[tuple(map(slice, array.shape))]
# if you don't need full ND generality: in 3D this can be written
# much more readable as
# m, n, k = array.shape
# idx = np.ogrid[:m, :n, :k]

idx[a] = order

print(np.all(array[idx] == np.sort(array, axis=a)))

Выход:

True

Объяснение: Мы должны указать для каждого элемента выходного массива полный индекс соответствующего элемента входного массива. Таким образом, каждый индекс во входном массиве имеет ту же форму, что и выходной массив, или должен транслироваться в эту форму.

Индексы для осей, по которым мы не сортируем / не сортируем, остаются на месте. Поэтому нам нужно передать широковещательный диапазон (array.shape[i]) для каждого из них. Самый простой способ - использовать ogrid для создания такого диапазона для всех измерений (если бы мы использовали это напрямую, массив возвращался без изменений.), А затем заменил индекс, соответствующий g для оси сортировки, выводом argsort,

Источник

Ответ @Hameer работает, хотя он может использовать некоторые упрощения и объяснения.

sort а также argsort работаем на последней оси. argsort возвращает трехмерный массив той же формы, что и оригинал. Значения - это индексы на этой последней оси.

In [17]: np.argsort(arr, axis=2)
Out[17]: 
array([[[1, 0],
        [1, 0]],

       [[0, 1],
        [1, 0]]], dtype=int32)
In [18]: _.shape
Out[18]: (2, 2, 2)
In [19]: idx=np.argsort(arr, axis=2)

Чтобы использовать это, нам нужно построить индексы для других измерений, которые транслируются в той же (2,2,2) форме. ix_ это удобный инструмент для этого.

Просто используя idx как один из ix_ входы не работают:

In [20]: np.ix_(range(2),range(2),idx)
....
ValueError: Cross index must be 1 dimensional

Вместо этого я использую последний диапазон, а затем игнорирую его. @Hameer вместо этого создает 2d ix_, а затем расширяет их.

In [21]: I,J,K=np.ix_(range(2),range(2),range(2))
In [22]: arr[I,J,idx]
Out[22]: 
array([[[ 0.62078744,  0.81774634],
        [ 0.29718462,  0.43912609]],

       [[ 0.1266578 ,  0.82282054],
        [ 0.79134389,  0.98180375]]])

Таким образом, индексы для других измерений работают с (2,2,2) массив idx:

In [24]: I.shape
Out[24]: (2, 1, 1)
In [25]: J.shape
Out[25]: (1, 2, 1)

Это основы для построения других индексов, когда вам дается многомерный индекс для одного измерения.

@Paul строит те же индексы с ogrid:

In [26]: np.ogrid[slice(2),slice(2),slice(2)]  # np.ogrid[:2,:2,:2]
Out[26]: 
[array([[[0]],

        [[1]]]), array([[[0],
         [1]]]), array([[[0, 1]]])]
In [27]: _[0].shape
Out[27]: (2, 1, 1)

ogrid как class работает с ломтиками, а ix_ требует список / массив / диапазон.

argsort для многомерного ndarray (с 2015 года) работает с 2d массивом, но применяется та же логика (найдите индекс (ы) диапазона, который вещает с argsort).

Источник

Вот векторизованная реализация. Он должен быть N-мерным и немного быстрее, чем вы делаете.

import numpy as np


def sort1(array, args):
    array_sort = np.zeros_like(array)
    for i in range(array.shape[0]):
        for j in range(array.shape[1]):
            array_sort[i, j] = array[i, j, args[i, j]]

    return array_sort


def sort2(array, args):
    shape = array.shape
    idx = np.ix_(*tuple(np.arange(l) for l in shape[:-1]))
    idx = tuple(ar[..., None] for ar in idx)
    array_sorted = array[idx + (args,)]

    return array_sorted


if __name__ == '__main__':
    array = np.random.rand(5, 6, 7)
    idx = np.argsort(array)

    result1 = sort1(array, idx)
    result2 = sort2(array, idx)

    print(np.array_equal(result1, result2))
Источник
Другие вопросы по тегам