python - 交错两个 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(对于 a)变为 1(对于 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 微秒，尽管您的情况可能会有所不同。