Julia 中的高效部分置换排序

Efficient partial permutation sort in Julia

我正在处理一个问题,该问题需要在 Julia 中按量级进行部分排列排序。如果 x 是维数 p 的向量,那么我需要的是第一个 k 索引对应于 xk 个分量,它们将首先出现在按 x 的绝对值进行部分排序。

参考 Julia 的排序函数 here. Basically, I want a cross between sortperm and select!. When Julia 0.4 is released, I will be able to obtain the same answer by applying sortperm! (this function) 到索引向量并选择其中的第一个 k。但是,在这里使用 sortperm! 并不理想,因为它会对 x 的剩余 p-k 索引进行排序,而我不需要这些索引。

执行部分置换排序最节省内存的方法是什么?我通过查看 sortperm 源代码破解了一个解决方案。但是,由于我不熟悉 Julia 在那里使用的排序模块,我不确定我的方法是否智能。

一个重要的细节:我可以忽略这里的重复或歧义。换句话说,我不关心 abs() 两个分量 2-2 的索引排序。我的实际代码使用浮点值,因此出于实际目的永远不会出现完全相等的情况。

# initialize a vector for testing
x  = [-3,-2,4,1,0,-1]
x2 = copy(x)
k  = 3    # num components desired in partial sort
p  = 6    # num components in x, x2

# what are the indices that sort x by magnitude?
indices = sortperm(x, by = abs, rev = true)

# now perform partial sort on x2
select!(x2, k, by = abs, rev = true)

# check if first k components are sorted here
# should evaluate to "true"
isequal(x2[1:k], x[indices[1:k]])

# now try my partial permutation sort
# I only need indices2[1:k] at end of day!
indices2 = [1:p]
select!(indices2, 1:k, 1, p, Base.Perm(Base.ord(isless, abs, true, Base.Forward), x))

# same result? should evaluate to "true"
isequal(indices2[1:k], indices[1:k])

编辑:使用建议的代码,我们可以简要比较更大向量上的性能:

p = 10000; k = 100;    # asking for largest 1% of components
x  = randn(p); x2 = copy(x);

# run following code twice for proper timing results
@time {indices = sortperm(x, by = abs, rev = true); indices[1:k]};
@time {indices2 = [1:p]; select!(indices2, 1:k, 1, p, Base.Perm(Base.ord(isless, abs, true, Base.Forward), x))};
@time selectperm(x,k);

我的输出:

elapsed time: 0.048876901 seconds (19792096 bytes allocated)
elapsed time: 0.007016534 seconds (2203688 bytes allocated)
elapsed time: 0.004471847 seconds (1657808 bytes allocated)

以下版本似乎相对 space 高效,因为它仅使用与输入数组长度相同的整数数组:

function selectperm (x,k)
    if k > 1 then
        kk = 1:k
    else
        kk = 1
    end
    z = collect(1:length(x))
    return select!(z,1:k,by = (i)->abs(x[i]), rev = true)
end    
x  = [-3,-2,4,1,0,-1]

k  = 3    # num components desired in partial sort
print (selectperm(x,k))

输出为:

[3,1,2]

...正如预期的那样。

我不确定它是否比最初提出的解决方案使用更少的内存(尽管我怀疑内存使用情况相似)但代码可能更清晰并且 确实如此仅生成第一个 k 个索引,而原始解决方案生成所有 p 个索引。

(编辑)

selectperm() 已被编辑以处理如果 k=1 在对 select!().

的调用中出现的 BoundsError