tl; dr

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软件包，当对我的主数组进行 sorting 时，它可以将numpy吹倒。不幸的是，我需要获取结果索引，以便也可以转换非主数组。

因此：与使用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)

因此，您最初尝试使用连接并随后进行合并排序实际上比我编写的直接利用预排序功能的代码要快。关于SO的问题类似asked before，并且产生了类似的基准。查看merge sort algorithm的详细信息，我认为这可能是由于合并两个排序的列表与合并排序的最佳情况性能场景相距仅一步之遥。

高效地对两个预排序的numpy数组进行argsort

1 个答案:

tl; dr

没有合并排序的方法

测试/计时