排序数组并返回已排序数组的原始索引

Question

给定一个numpy数组,如何在其中找到索引序列,以便对结果进行排序？

例如,给定x=[4,2,6],结果将是[1,0,2],因为[x[1],x[0],x[2]]是排序的.

我知道有很多可用的Python函数argsort()可以完成这项工作,但我需要自己实现这个排序功能.有什么建议？

Answer 1

首先,您可以使用enumerate将任何可迭代的值转换为(索引,值)对的可迭代.

但是,如果你只是对它们进行排序,它将按索引排序,这不是很有用.您希望按每个(索引,值)对中的值进行排序.通常,在Python中,您可以通过传递一个键函数来实现sorted.如该文档中的示例所示,itemgetter这里提供了完美的键功能.您可以轻松修改自定义排序功能,以同样的方式使用键功能sorted,虽然如果没有看到自定义排序功能,有点难以告诉您如何做到这一点.¹

但在这种情况下,您可以使用Decorate-Sort-Undecorate成语.您只想按每个(索引,值)对中的值进行排序,因此您只需要对"装饰"进行反向操作即可.并且,如果您只希望索引而不是值"未整理",则只需删除值即可.

所以:

indexed = enumerate(arr)
decorated = ((value, index) for index, value in indexed)
sortedpairs = my_sort_function(decorated)
indices = np.fromiter(index for (value, index) in sortedpairs)

...或者,把它们放在一起:

sortedpairs = my_sort_function((value, index) for index, value in enumerate(arr))
indices = np.fromiter(index for (value, index) in sortedpairs)

(当然你可以把它作为一个单行,但我认为这两行是最好的可读性平衡.)

如果您不允许使用均匀enumerate,这是用您自己的功能替换的最简单的内置函数之一.事实上,文档甚至会告诉你如何做到这一点:

def my_enumerate(sequence, start=0):
    n = start
    for elem in sequence:
        yield n, elem
        n += 1

或者,因为您不需要自定义起始值:

def my_enumerate(sequence):
    n = 0
    for elem in sequence:
        yield n, elem
        n += 1

但是现在,你是否可以做同样的事情,同时仍然采取(至少一些)numpy的优势,保持一切为数组而不是使用iterables？

当然.我们可以做同样的事情enumerate,甚至把值放在底部,所以我们不需要整个翻转步骤:

decorated = np.stack((arr, np.arange(len(arr))))

...然后排序.我假设您的自定义排序功能对列进行排序.也许你需要传递一个axis论点,或者排序decorated.T,或者其他什么; 你应该知道你自己的函数的API.

sorted_pairs = my_sorted_array_function(decorated)

现在,我们只需要索引行:

indices = sorted_pairs[1]

_{1.对于初始实现,只需将每个更改x < y为key(x) < key(y),并使其正常工作.然后,您可以通过缓存键值来弄清楚如何优化它,这样您key每个元素只调用一次而不是每个元素调用一次log(N).}