交错两个numpy索引数组,每个数组一个项目

Question

我有两个有序的numpy数组,我想交错它们,以便我从第一个数组中取出一个项目,然后从第二个数据中获取另一个项目,然后返回到第一个 - 获取下一个项目大于我刚刚从中获取的项目第二个等等.这些实际上是其他数组的索引数组,只要操作是向量化的,我就可以对原始数组进行操作(但当然,作为向量操作处理索引数组会非常棒).

示例(可以假设数组的交集为空)

a = array([1,2,3,4,7,8,9,10,17])
b = array([5,6,13,14,15,19,21,23])

我想得到[1,5,7,13,17,19]

Answer 1

矢量化解决方案(教学风格,易于理解)

我们可以通过使用鉴别器索引扩充数组来进行矢量化,例如a标记0并b标记1:

a_t = np.vstack((a, np.zeros_like(a)))
b_t = np.vstack((b, np.ones_like(b)))

现在,让我们结合并排序:

c = np.hstack((a_t, b_t))[:, np.argsort(np.hstack((a, b)))]
array([[ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 13, 14, 15, 17, 19, 21, 23],
       [ 0,  0,  0,  0,  1,  1,  0,  0,  0,  0,  1,  1,  1,  0,  1,  1,  1]])

您可以看到现在元素按顺序排列但保留了它们的标记,因此我们可以看到哪些元素来自a和来自哪些元素b.

所以,让我们选择第一个元素和标签从0(for a)到1(for b)再返回的每个元素:

c[:, np.concatenate(([True], c[1, 1:] != c[1, :-1]))][0]
array([ 1,  5,  7, 13, 17, 19])

高效的矢量化解决方案

您可以通过将项目及其标记保存在单独(但并行)的数组中来更有效地执行此操作:

ab = np.hstack((a, b))
s = np.argsort(ab)
t = np.hstack((np.zeros_like(a), np.ones_like(b)))[s]
ab[s][np.concatenate(([True], t[1:] != t[:-1]))]
array([ 1,  5,  7, 13, 17, 19])

这比上述解决方案稍微有效一些; 虽然您的条件可能会有所不同,但平均为45而不是90微秒.