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

Posted 2023-02-23

【中文标题】使用 numpy 从过滤后的排序数组返回索引

【英文标题】：return indices from filtered, sorted array with numpy

【发布时间】：2015-12-28 15:54:09

【问题描述】：

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

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

【问题讨论】：

是否有任何解决方案适合您？

【参考方案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])

【参考方案2】：

掩码数组很简洁（但可能不是特别有效）

x = np.random.rand(4);

inverse_mask = x < 0.5
m_x = np.ma.array(x, mask=np.logical_not(inverse_mask))
sorted_indeces = m_x.argsort(fill_value=1)
filtered_sorted_indeces = sorted_indeces[:np.sum(inverse_mask)]

【参考方案3】：

一种解决方案：

已创建排序索引数组 (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)