高效的numpy argsort与条件同时保持原始指数

Question

我想知道在给定条件的情况下,在保留原始索引的同时,最有效的方法是做一个数组的argsort

import numpy as np

x = np.array([0.63, 0.5, 0.7, 0.65])

np.argsort(x)
#Corrected argsort(x) solution
Out[99]: array([1, 0, 3, 2])

我想在x> 0.6的条件下对这个数组进行argsort.由于0.5 <0.6,因此不应包括指数1.

x = np.array([0.63, 0.5, 0.7, 0.65])
index = x.argsort()
list(filter(lambda i: x[i] > 0.6, index))

[0,3,2]

这是低效的,因为它没有矢量化.

编辑:过滤器将消除大多数元素.理想情况下,它首先进行过滤,然后进行排序,同时保留原始索引.

Answer 1

方法1(与Tai的方法相同但使用整数索引)

太晚了,如果我的解决方案是重复已经发布的解决方案 - ping我,我将删除它.

def meth_agn_v1(x, thresh):
    idx = np.arange(x.size)[x > thresh]
    return idx[np.argsort(x[idx])]

然后,

In [143]: meth_agn_v1(x, 0.5)
Out[143]: array([0, 3, 2])

方法2(显着的性能改进)

这使用了我的答案的最后一部分(与Tai的方法比较)中表达的相同的想法,即整数索引比布尔索引更快(对于要选择的少量预期元素)并且完全避免创建初始索引.

def meth_agn_v2(x, thresh):
    idx, = np.where(x > thresh)
    return idx[np.argsort(x[idx])]

定时

In [144]: x = np.random.rand(100000)

In [145]: timeit meth_jp(x, 0.99)
100 loops, best of 3: 7.43 ms per loop

In [146]: timeit meth_alex(x, 0.99)
1000 loops, best of 3: 498 µs per loop

In [147]: timeit meth_tai(x, 0.99)
1000 loops, best of 3: 298 µs per loop

In [148]: timeit meth_agn_v1(x, 0.99)
1000 loops, best of 3: 232 µs per loop

In [161]: timeit meth_agn_v2(x, 0.99)
10000 loops, best of 3: 95 µs per loop

v1与Tai方法的比较

我答案的第一个版本与Tai的答案非常相似,但不完全相同.

Tai的方法最初发布:

def meth_tai(x, thresh):
    y = np.arange(x.shape[0])
    y = y [x > thresh]  
    x = x [x > thresh] # x = x[y] is used in my method
    y[np.argsort(x)]

因此,我的方法与使用整数数组索引而不是Tai使用的布尔索引不同.对于少数选定的元素,整数索引比布尔索引更快,这使得这种方法比Tai的方法更有效,即使在Tai优化了他的代码之后.

Answer 2

来参加派对有点晚了.我们的想法是,我们可以根据另一个数组的排序索引对数组进行排序.

y = np.arange(x.shape[0]) # y for preserving the indices
mask = x > thresh
y = y[mask]  
x = x[mask]
ans = y[np.argsort(x)]    # change order of y based on sorted indices of x

该方法是添加一个y仅用于记录索引的数组x.然后我们根据布尔索引过滤掉两个数组x > thresh.然后,排序x与argsort.最后,使用索引返回argsort来改变顺序y!