使用 numpy 从过滤、排序的数组中返回索引

Question

在 python/numpy 中执行以下操作的最直接方法是什么？

• 以随机数组开始 x
• 过滤掉元素 x < .5
• 按大小对剩余值进行排序
• 返回x对应于这些值的（原始）索引

Answer 1

找到的面具x < 0.5，并x.argsort()在这里似乎是强制性的。一旦你有了这两个，你就可以使用排序索引对掩码数组进行排序，并在排序索引上使用这个掩码来取回与满足掩码条件的排序索引相对应的索引。因此，您将再添加一行代码，如下所示 -

mask = x < 0.5
sort_idx = x.argsort()
out = sort_idx[mask[sort_idx]]

示例分步运行 -

In [56]: x
Out[56]: array([ 0.8974009 ,  0.30127187,  0.71187137,  0.04041124])

In [57]: mask
Out[57]: array([False,  True, False,  True], dtype=bool)

In [58]: sort_idx
Out[58]: array([3, 1, 2, 0])

In [59]: mask[sort_idx]
Out[59]: array([ True,  True, False, False], dtype=bool)

In [60]: sort_idx[mask[sort_idx]]
Out[60]: array([3, 1])

Answer 2

一种解决方案：

• 创建排序索引数组（argsort）
• x为排序小于阈值创建掩码
• 将掩码应用于排序索引数组

例子：

import numpy as np

# x = np.random.rand(4)
x = np.array([0.96924269, 0.30592608, 0.03338015, 0.64815553])
solution = np.array([2, 1])

sorted_idx = np.argsort(x)
idx_mask = (x[sorted_idx] < 0.5)
sorted_filtered_idx = sorted_idx[idx_mask]

assert np.all(sorted_filtered_idx == solution)