Performanly argsort два предварительно отсортированных массива numpy

Question

У меня есть несколько групп массивов.Внутри каждой группы все массивы одномерные, все одинаковой длины.Внутри каждой группы есть один первичный массив, который уже отсортирован.

Например:

grp_1 = [
    np.array([10, 20, 30, 40]),
    np.array(["A", "C", "E", "G"]),
    ]

grp_2 = [
    np.array([15, 25, 35]),
    np.array(["Z", "Y", "X"]),
    ]

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

def combine_groups(groups):
    combined_arrays = [np.concatenate([grp[idx] for grp in groups]) for idx in range(len(groups[0]))]
    sort_indices = np.argsort(combined_arrays[0], kind="mergesort")
    # Merge sort rather than quicksort because the former is stable
    return [arr[sort_indices] for arr in combined_arrays]

Это работает и прекрасно, но (для массивов, гораздо больших, чем в этом примере), это намного медленнее, чем нужно.Сортировка слиянием - это O (N log (N)), тогда как объединяющиеся массивы , которые уже отсортированы , должны быть делом O (N).

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

Итак: вышеприведенное возможно быстрее, чем использование numpy's argsort?

Answer 1

tl; dr

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

Методы без сортировки слиянием

Просто заархивируйте ваши группы и затем используйте cytoolz.merge_sorted:

from cytoolz import merge_sorted

# it will be an iterator that yields (10, 'A'), (15, 'Z'), (20, 'C'), (25, 'Y'), (30, 'E'), (35, 'X'), (40, 'G')
it = merge_sorted(zip(*grp_1), zip(*grp_2))

# unzip the tuples back into your desired arrays
grp_comb = [np.array(d) for d in zip(*it)]
print(grp_comb)

Вывод:

[array([10, 15, 20, 25, 30, 35, 40]), array(['A', 'Z', 'C', 'Y', 'E', 'X', 'G'], dtype='<U1')]

В качестве альтернативы, есливы действительно хотите объединить свои группы с помощью косвенной сортировки, например numpy.argsort, вы можете использовать ndarray.searchsorted:

ix = grp_1[0].searchsorted(grp_2[0])
grp_comb= [np.insert(grp_1[i], ix, grp_2[i]) for i in range(2)]
print(grp_comb)

Выход:

[array([10, 15, 20, 25, 30, 35, 40]), array(['A', 'Z', 'C', 'Y', 'E', 'X', 'G'], dtype='<U1')]

Тестирование / синхронизация

Я использовал следующий код, чтобы проверить, выводит ли мои ответы вывод, идентичный отправленной вами функции combine_groups, и приурочить различные методы:

from cytoolz import merge_sorted
import numpy as np
from numpy.testing import assert_array_equal

grp_1 = [
    np.array([10, 20, 30, 40]),
    np.array(["A", "C", "E", "G"]),
    ]

grp_2 = [
    np.array([15, 25, 35]),
    np.array(["Z", "Y", "X"]),
    ]

def combine_groups(*groups):
    combined_arrays = [np.concatenate([grp[idx] for grp in groups]) for idx in range(len(groups[0]))]
    sort_indices = np.argsort(combined_arrays[0], kind="mergesort")
    # Merge sort rather than quicksort because the former is stable
    return [arr[sort_indices] for arr in combined_arrays]

def combine_groups_ms(*groups):
    it = merge_sorted(*(zip(*g) for g in groups))
    return [np.array(d) for d in zip(*it)]

def combine_groups_ssi(g0, g1):
    ix = g0[0].searchsorted(g1[0])
    return [np.insert(g0[i], ix, g1[i]) for i in range(len(g0))]

expected = combine_groups(grp_1, grp_2)
assert_array_equal(combine_groups_ms(grp_1, grp_2), expected)
assert_array_equal(combine_groups_ssi(grp_1, grp_2), expected)

Вот время:

%%timeit
combine_groups(grp_1, grp_2)
6.84 µs ± 154 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%%timeit
combine_groups_ms(grp_1, grp_2)
10.4 µs ± 249 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%%timeit
combine_groups_ssi(grp_1, grp_2)
36.3 µs ± 1.27 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

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