python - 根据更改阈值过滤 np 数组中的值的最快方法

Question

我想过滤一个数组arr基于一些阈值。

arr = np.array([2,2,2,2,2,5,5,5,1])
thresholds = np.array([4,1])

我要过滤arr基于 thresholds 中的值当arr中的值大于阈值

我的想法是为每个阈值创建一个掩码

预期结果:

# [[False False False False False  True  True  True False]
#  [ True  True  True  True  True  True  True  True False]]

在 Python 中执行此操作的一种方法:

mask = [True if x>condi else False for condi in thresholds for x in arr]
mask = np.reshape(mask,(2,9))

然后通过 filteredarr = arr[mask[i]] 获取过滤后的数组哪里i为相关阈值的索引

有没有更好的方法(性能方面)在 Python 中做到这一点？特别是我正在处理大数组(arr 的 len 约为 250000，thresholds 还没有特定的 len，但我期待一个大数组)？

编辑:数据的最终预期输出为 [array([5, 5, 5]), array([2, 2, 2, 2, 2, 5, 5, 5])]

Answer 1

可以使用以下方法轻松获得面具

mask = arr[None,:]>thresholds[:,None]
mask

# Output
# array([[False, False, False, False, False,  True,  True,  True, False],
#        [ True,  True,  True,  True,  True,  True,  True,  True, False]], dtype=bool)

这个想法是通过使用 None 添加额外的轴(与 np.newaxis 作用相同)来放大维数，然后比较数组元素 -明智的。

一旦我们有了掩码，我们就可以使用各种方法过滤数据，其中选择很大程度上取决于您的问题:

• 当然可以

  res = [arr[m] for m in mask]
  # [array([5, 5, 5]), array([2, 2, 2, 2, 2, 5, 5, 5])]
  

  为了获取过滤后的数据的列表，但总体来说比较慢。

• 如果您需要进一步进行数值计算，我将创建一个掩码数组，其中仅考虑过滤后的数据:

  m = np.zeros_like(mask).astype(np.int)
  m[:] = arr
  res = np.ma.masked_where(~mask,m)
  

  现在每行对应于根据相应阈值过滤后的数据。屏蔽数组允许您继续使用许多函数，例如 mean 或 std

  res.mean(axis=1)
  # masked_array(data = [5.0 3.125],
  #             mask = [False False],
  #       fill_value = 1e+20)
  
  res.mean(axis=1).compressed()
  # array([ 5.   ,  3.125])