有效地重新排列 NumPy 数组

Question

假设我有一个简单的一维 NumPy 数组：

x = np.random.rand(1000)

然后我检索排序后的索引：

idx = np.argsort(x)

但是，我需要将索引列表移动到 idx 的前面。因此，假设indices = [10, 20, 30, 40, 50] 必须始终是前 5 个，然后其余的将从 idx 开始（减去在 indices 中找到的索引）

实现这一点的一种天真的方法是：

indices = np.array([10, 20, 30, 40, 50])
out = np.empty(idx.shape[0], dtype=int64)
out[:indices.shape[0]] = indices

n = indices.shape[0]
for i in range(idx.shape[0]):
    if idx[i] not in indices:
        out[n] = idx[i] 
        n += 1

有没有办法有效地做到这一点，并且可能就地完成？

Answer 1

您可以构建一个掩码，其中indices 包含在idx 和np.in1d 中，并且只需连接两个索引数组：

m = np.in1d(idx, indices)
out = np.r_[indices, idx[~m]]

Answer 2

方法#1

一种方法是使用 np.isin 屏蔽 -

mask = np.isin(idx, indices, invert=True)
out = np.r_[indices, idx[mask]]

方法#2：跳过第一个argsort

另一个使这些给定索引最小化，从而迫使它们以argsorting 开头。对于这种方法，我们不需要 idx，因为无论如何我们都在我们的解决方案中进行 argsort-ing -

def argsort_constrained(x, indices):
    xc = x.copy()
    xc[indices] = x.min()-np.arange(len(indices),0,-1)
    return xc.argsort()

基准测试 - 仔细观察

让我们研究一下跳过启动argsortidx 的计算这整个事情如何帮助我们使用第二种方法。

我们将从给定的示例开始：

In [206]: x = np.random.rand(1000)

In [207]: indices = np.array([10, 20, 30, 40, 50])

In [208]: %timeit argsort_constrained(x, indices)
38.6 µs ± 1.39 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [209]: idx = np.argsort(x)

In [211]: %timeit np.argsort(x)
27.7 µs ± 122 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [212]: %timeit in1d_masking(x, idx, indices)
     ...: %timeit isin_masking(x, idx, indices)
44.4 µs ± 421 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
50.7 µs ± 303 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

请注意，如果您在这些小型数据集上使用 np.concatenate 代替 np.r_，您可以做得更好。

因此，argsort_constrained 的总运行时间成本约为 38.6 µs，而其他两个带有掩码的运行时间成本约为 27.7 µs，在它们各自的计时数字之上。

让我们通过10x 扩大所有内容并进行相同的实验：

In [213]: x = np.random.rand(10000)

In [214]: indices = np.sort(np.random.choice(len(x), 50, replace=False))

In [215]: %timeit argsort_constrained(x, indices)
740 µs ± 3.13 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [216]: idx = np.argsort(x)

In [217]: %timeit np.argsort(x)
731 µs ± 14.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [218]: %timeit in1d_masking(x, idx, indices)
     ...: %timeit isin_masking(x, idx, indices)
1.07 ms ± 47.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
1.02 ms ± 4.02 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

同样，使用掩码的单个运行时成本高于使用argsort_constrained。随着我们进一步扩大规模，这种趋势应该会继续下去。