在迭代部分更新的数组中重复查找 K 个最大值索引的最快方法

我尝试实现基于 C++ 容器的 Cython 解决方案（用于 64 位浮点值）。好消息是它比天真的要快。坏消息是它非常复杂，而且速度并不快：快 3~4 倍。np.argpartition

一个主要问题是 Cython 没有实现最有用的容器。可以使用类型实现此容器，但它会使代码更加复杂，并且效率也较低（由于内存中额外的缓存不友好的间接）。如果可以保证 中没有重复项，那么性能会明显更好（最多 x2），因为可以改用。此外，Cython 似乎不接受最近的 C++11/C++17/C++20 功能，使代码读/写起来更加繁琐。这很可悲，因为 [一些功能，如和 rvalues-references] 可以使代码更快。std::multimapstd::map<Key, std::vector<Value>>bstd::mapextract

另一个主要问题是执行时间受缓存未命中（在我的机器上为 >75%）的限制，因为二进制 RB 树对缓存不友好。问题是整体数据结构很可能比 CPU 缓存大。事实上，至少需要存储键值，更不用说存储树数据结构的节点指针需要类似数量的内存，并且大多数处理器缓存都小于 30 MB。这主要是由于随机访问而导致的更新期间的一个问题：每次查找/插入都需要在 RAM 中获取，并且 RAM 的延迟通常为几十纳秒。此外，（C++） RB 树不支持密钥更新，因此需要删除+插入。我试图使用并行预取方法来缓解这个问题。不幸的是，它在实践中通常较慢......750_000*(8*2+4) = 15_000_000 byteslog2(nsel)

在实践中，K 最大项的提取速度非常快（树中的 1000 个项和 750_000 个值大约需要几微秒），而更新大约需要 1.0-1.5 毫秒。同时，需要 ~4.5 毫秒。np.argpartition

一些人报告了（例如。这里）当项目数量相当大时，这实际上是相当慢的。因此，使用另一个非标准 C++ 实现可能是个好主意。在这种情况下，我希望 B 树会更快。Google Abseil 库包含这样的容器，它们的速度肯定要快得多。话虽如此，它肯定需要包装一些代码，这可能很乏味。或者，可以编写一个完整的 C++ 类并从 Cython 调用它。std::map

实现

下面是实现（以及最后的用法示例）：

• maxtree.pyx：

# distutils: language = c++

import numpy as np
cimport numpy as np
cimport cython

# See: https://cython.readthedocs.io/en/latest/src/userguide/wrapping_CPlusPlus.html
from libcpp.vector cimport vector
from libcpp.map cimport map
from libcpp.pair cimport pair
from cython.operator cimport dereference as deref, preincrement as inc


@cython.boundscheck(False)  # Deactivate bounds checking
@cython.wraparound(False)   # Deactivate negative indexing
cdef class MaxTree:
    cdef map[double, vector[int]] data
    cdef int itemCount

    # Build a tree from `b`
    def __init__(self, double[::1] b):
        cdef map[double, vector[int]].iterator it
        cdef pair[double, vector[int]] node
        cdef double val
        cdef int i

        # Temporary node used to ease insertion
        node.second.resize(1)

        # Iterate over `b` items so to add them in the tree
        for i in range(b.size):
            val = b[i]
            it = self.data.find(val)

            if it == self.data.end():
                # Value not found: add a new node
                node.first = val
                node.second[0] = i
                self.data.insert(node)
            else:
                # Value found: adds a new duplicate in an existing node
                deref(it).second.push_back(i)

        self.itemCount = b.size

    def size(self):
        return self.itemCount

    # Get the index (in the original `b` array) of the K-largest values
    def getKlargest(self, int count):
        cdef map[double, vector[int]].reverse_iterator rit
        cdef int vecSize
        cdef int* vecData
        cdef int i, j
        cdef int[::1] resultView

        if count > self.itemCount:
            count = self.itemCount

        result = np.empty(count, dtype=np.int32)
        resultView = result
        i = 0

        rit = self.data.rbegin()
        while rit != self.data.rend():
            vecSize = deref(rit).second.size()
            vecData = deref(rit).second.data()
            # Note: indices are not always sorted here due to the update
            for j in range(vecSize-1, -1, -1):
                resultView[i] = vecData[j]
                i += 1
                count -= 1
                if count <= 0:
                    return resultView
            inc(rit)

        return result

    # Set the values of `b` at the index `index` to `values` and update the tree accordingly
    def update(self, double[::1] b, int[::1] index, double[::1] values):
        cdef map[double, vector[int]].iterator it
        cdef pair[double, vector[int]] node
        #cdef pair[map[double, vector[int]].iterator, bool] infos
        cdef int idx, i, j, vecSize, indexSize
        cdef double oldValue, newValue
        cdef int* vecData

        assert b.size == self.itemCount
        assert index.size == values.size
        assert np.min(index) >= 0 and np.max(index) < b.size

        # Temporary node used to ease insertion
        node.second.resize(1)

        for i in range(index.size):
            idx = index[i]
            oldValue = b[idx]
            newValue = values[i]

            it = self.data.find(oldValue)
            assert it != self.data.end()

            # Update the tree
            if deref(it).second.size() == 1:
                # Remove the node from the tree and add a new one because keys are immutable
                # Assume `index` is correct/coherent and the tree is correctly updated for sake of performance
                #assert deref(it).second[0] == idx
                self.data.erase(it)
                node.first = newValue
                node.second[0] = idx
                infos = self.data.insert(node)
                inserted = infos.second
                if not inserted:
                    # Duplicate
                    it = infos.first
                    deref(it).second.push_back(idx)
            else:
                # Tricky case due to duplicates (untested)
                vecData = deref(it).second.data()
                vecSize = deref(it).second.size()
                # Search the element and remove it
                for j in range(vecSize):
                    if vecData[j] == idx:
                        vecData[j] = vecData[vecSize-1]
                        deref(it).second.pop_back()
                        break

            # Update `b`
            b[idx] = values[i]

• setup.py：

# setup.py

from setuptools import setup
from Cython.Build import cythonize

setup(ext_modules=cythonize("maxtree.pyx"))

• main.py：

# Usage:

import numpy as np
import maxtree
np.random.seed(0)
b = np.random.rand(750_000)
nchange = 1_000
ichange = np.random.randint(0, b.size, nchange).astype(np.int32)

tree = maxtree.MaxTree(b)
tree.getKlargest(nchange)
tree.update(b, ichange, b[ichange]*0.999)

• 要运行的命令：python3 setup.py build_ext --inplace -q