如何在 Python 中高效且增量地对向量进行 argsort 排序？

Question

在收集一些样本的循环中，我需要不时地获取一些关于它们的排序索引的统计信息，argsort 正好返回我需要的信息。但是，每次迭代只添加一个样本，并且不断将整个样本数组传递给 argsort 函数是一种巨大的资源浪费，特别是因为样本数组非常庞大。难道没有等效于 argsort 的增量高效技术吗？

我相信一个有效的增量 argsort 函数可以通过维护一个有序的样本列表来实现，一旦新样本到达，该列表可以是 searched 用于正确的 insertion 索引。然后可以利用这些索引来维护样本列表的顺序以及生成增量的类似 argsort 的所需输出。 到目前为止，我已经使用了@Divakar 的searchsorted2d 函数，稍作修改以获得插入索引，并构建了一些例程，如果在每次插入样本后调用它可以获得所需的输出（b = 1）。 然而，这是低效的，我想在收集第 k 个样本后调用例程（例如b = 10）。 在批量插入的情况下，searchsorted2d 似乎返回了不正确的索引，这就是我停止了！

import time
import numpy as np

# By Divakar
# See https://stackoverflow.com/a/40588862
def searchsorted2d(a, b):
    m, n = a.shape
    max_num = np.maximum(a.max() - a.min(), b.max() - b.min()) + 1
    r = max_num * np.arange(m)[:,np.newaxis]
    p = np.searchsorted((a + r).ravel(), (b + r).ravel()).reshape(b.shape)
    return p #- n * (np.arange(m)[:,np.newaxis])

# The following works with batch size b = 1,
# but that is not efficient ...
# Can we make it work for any b > 0 value?
class incremental(object):
    def __init__(self, shape):
        # Express each row offset
        self.ranks_offset = np.tile(np.arange(shape[1]).reshape(1, -1),
                                    (shape[0], 1))
        # Storage for sorted samples
        self.a_sorted = np.empty((shape[0], 0))
        # Storage for sort indices
        self.a_ranks = np.empty((shape[0], 0), np.int)

    def argsort(self, a):
        if self.a_sorted.shape[1] == 0: # Use np.argsort for initialization
            self.a_ranks = a.argsort(axis=1)
            self.a_sorted = np.take_along_axis(a, self.a_ranks, 1)
        else: # In later itterations,
            # searchsorted the input increment
            indices = searchsorted2d(self.a_sorted, a)
            # insert the stack pos to track the sorting indices
            self.a_ranks = np.insert(self.a_ranks, indices.ravel(),
                                     self.ranks_offset.ravel() +
                                     self.a_ranks.shape[1]).reshape((n, -1))
            # insert the increments to maintain a sorted input array
            self.a_sorted = np.insert(self.a_sorted, indices.ravel(),
                                      a.ravel()).reshape((n, -1))
        return self.a_ranks

M = 1000 # number of iterations
n = 16   # vector size
b = 10   # vectors batch size

# Storage for samples
samples = np.zeros((n, M)) * np.nan

# The proposed approach
inc = incremental((n, b))

c = 0 # iterations counter
tick = time.time()
while c < M:
    if c % b == 0: # Perform batch computations
        #sample_ranks = samples[:,:c].argsort(axis=1)
        sample_ranks = inc.argsort(samples[:,max(0,c-b):c]) # Incremental argsort

        ######################################################
        # Utilize sample_ranks in some magic statistics here #
        ######################################################

    samples[:,c] = np.random.rand(n) # collect a sample
    c += 1 # increment the counter
tock = time.time()

last = ((c-1) // b) * b
sample_ranks_GT = samples[:,:last].argsort(axis=1) # Ground truth
print('Compatibility: {0:.1f}%'.format(
      100 * np.count_nonzero(sample_ranks == sample_ranks_GT) / sample_ranks.size))
print('Elapsed time: {0:.1f}ms'.format(
      (tock - tick) * 1000))

我希望与 argsort 函数 100% 兼容，但它需要比调用 argsort 更有效。至于增量方法的执行时间，对于给定的示例，似乎 15 毫秒左右应该绰绰有余。 到目前为止，任何一种探索的技术都只能满足这两个条件中的一个。

---

长话短说，上面显示的算法似乎是order-statistic tree 的一种变体，用于估计数据等级，但是当批量添加样本时（b > 1）它无法这样做。到目前为止，它仅在一个一个插入样本时有效（b = 1）。但是，每次调用insert 时都会复制数组，这会导致巨大的开销并形成瓶颈，因此应批量添加样本而不是单独添加样本。

能否介绍一下更高效的增量argsort算法，或者至少弄清楚上面的如何支持批量插入（b > 1）？

如果你选择从我停止的地方开始，那么问题可以归结为修复以下快照中的错误：

import numpy as np

# By Divakar
# See https://stackoverflow.com/a/40588862
def searchsorted2d(a, b):
    m, n = a.shape
    max_num = np.maximum(a.max() - a.min(), b.max() - b.min()) + 1
    r = max_num * np.arange(m)[:,np.newaxis]
    p = np.searchsorted((a + r).ravel(), (b + r).ravel()).reshape(b.shape)
    # It seems the bug is around here...
    #return p - b.shape[0] * np.arange(b.shape[1])[np.newaxis]
    #return p - b.shape[1] * np.arange(b.shape[0])[:,np.newaxis]
    return p

n = 16  # vector size
b = 2   # vectors batch size

a = np.random.rand(n, 1) # Samples array
a_ranks = a.argsort(axis=1) # Initial ranks
a_sorted = np.take_along_axis(a, a_ranks, 1) # Initial sorted array

new_data = np.random.rand(n, b) # New block to append into the samples array
a = np.hstack((a, new_data)) #Append new block

indices = searchsorted2d(a_sorted, new_data) # Compute insertion indices
ranks_offset = np.tile(np.arange(b).reshape(1, -1), (a_ranks.shape[0], 1)) + a_ranks.shape[1] # Ranks to insert
a_ranks = np.insert(a_ranks, indices.ravel(), ranks_offset.ravel()).reshape((n, -1)) # Insert ransk according to their indices
a_ransk_GT = a.argsort(axis=1) # Ranks ground truth

mask = (a_ranks == a_ransk_GT)
print(mask) #Why they are not all True?
assert(np.all(mask)), 'Oops!' #This should not fail, but it does :(

---

似乎批量插入比我最初的想法更复杂，而且 searchsorted2d 不应该受到指责。以排序数组a = [ 1, 2, 5 ] 为例，要插入两个新元素block = [3, 4]。如果我们迭代并插入，那么np.searchsorted(a, block[i]) 将返回[2] 和[3]，这样就可以了。然而，如果我们调用np.searchsorted(a, block)（期望的行为——相当于没有插入的迭代），我们会得到[2, 2]。这对于实现增量argsort 是有问题的，因为即使np.searchsorted(a, block[::-1]) 也会导致相同的结果。任何想法？

Answer 1

事实证明，searchsorted 返回的索引不足以在处理批量输入时确保排序数组。如果正在插入的块包含两个乱序的条目，但它们最终会被相邻放置在目标数组中，那么它们将收到完全相同的插入索引，从而以当前顺序插入，从而导致故障。因此，输入块本身需要在插入之前进行排序。有关数字示例，请参见问题的最后一段。

通过对输入块进行排序，对剩余部分进行适配，得到了与argsort100.0%兼容的解决方案，而且非常高效（10×10块b = 10插入1000个条目耗时15.6ms） .这可以通过将问题中发现的错误 incremental 类替换为以下内容来重现：

# by Hamdi Sahloul
class incremental(object):
    def __init__(self, shape):
        # Storage for sorted samples
        self.a_sorted = np.empty((shape[0], 0))
        # Storage for sort indices
        self.a_ranks = np.empty((shape[0], 0), np.int)

    def argsort(self, block):
        # Compute the ranks of the input block
        block_ranks = block.argsort(axis=1)
        # Sort the block accordingly
        block_sorted = np.take_along_axis(block, block_ranks, 1)
        if self.a_sorted.shape[1] == 0: # Initalize using the block data
            self.a_ranks = block_ranks
            self.a_sorted = block_sorted
        else: # In later itterations,
            # searchsorted the input block
            indices = searchsorted2d(self.a_sorted, block_sorted)
            # update the global ranks
            self.a_ranks = np.insert(self.a_ranks, indices.ravel(),
                                     block_ranks.ravel() +
                                     self.a_ranks.shape[1]).reshape((block.shape[0], -1))
            # update the overall sorted array
            self.a_sorted = np.insert(self.a_sorted, indices.ravel(),
                                      block_sorted.ravel()).reshape((block.shape[0], -1))
        return self.a_ranks