python - Performantly 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 的 argsort 更快？

Answer 1

tl;dr

只需像您正在做的那样使用合并排序。上一个discussions和 benchmarks类似的问题都表明，如果您不自己编写一些 cython 代码(甚至可能不会)，您将无法击败您已经使用的方法。

没有归并排序的方法

只需压缩您的组，然后使用 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)

因此，您最初尝试使用连接然后进行合并排序实际上比我编写的直接利用预排序的代码更快。类似问题已经asked before在 SO 上，并产生了类似的基准。查看 merge sort algorithm 的详细信息，我认为这可能是因为合并两个排序列表距离合并排序的最佳案例性能场景仅一步之遥。