迭代部分更新数组中快速查找K个最大值索引的方案咨询

现有一个含≈750,000个元素的复数值数组a，需执行≈10^6次迭代，每次迭代更新≤1000个元素。每次迭代后，需基于a元素的平方绝对值构成的实数值数组b，找出K个最大值的索引（K≤50，通常≤10，索引无需排序）。

每次迭代更新的元素索引可视为随机，但必包含一个原最大值元素；更新后新的最大值可能来自未更新元素。当前使用b.argpartition(-K)[-K:]查找索引，但因b数组过大，全量遍历速度极慢，成为非线性反卷积算法CLEAN的性能瓶颈。此前针对K=1的场景，采用分块仅更新块最大值的方案实现了>7倍加速，但该方案无法直接扩展到K>1的情况。

问题1：如何实现二叉搜索树并高效查找K个最大值索引 #

Python标准库无内置平衡二叉搜索树，但可通过最大堆+哈希表组合模拟，或基于第三方平衡BST实现的工具类落地，以下是两种可行思路：

方法1：最大堆+哈希表（纯Python/Numba加速） #

核心逻辑是用堆维护候选最大值，同时用哈希表记录每个索引的最新b值，避免堆中存在过时的旧值：

• 初始化：先用numpy.argpartition找出初始前K个最大值索引，将对应的(-b[idx], idx)（用负数模拟最大堆）放入堆，同时用字典current_b记录每个索引的最新b值。
• 迭代更新：每次更新b中的元素时，将新的(-new_b_val, idx)推入堆，并更新current_b[idx] = new_b_val。
• 查找前K个最大值：从堆顶弹出元素，检查堆中存储的-val是否等于current_b[idx]（验证是否为最新值）：
  • 若相等，将该索引加入结果列表；
  • 若不相等，说明该元素是过时数据，直接丢弃；
  • 重复直到收集到K个有效索引，再将这些有效元素重新推回堆中（供下次使用）。

Numba优化示例：
纯Python的heapq在1e6次迭代下速度有限，可使用Numba实现自定义堆结构并JIT编译：

import numba as nb
import numpy as np

@nb.njit
def heapify(arr, n, i):
    largest = i
    l = 2 * i + 1
    r = 2 * i + 2
    if l < n and arr[l][0] > arr[largest][0]:
        largest = l
    if r < n and arr[r][0] > arr[largest][0]:
        largest = r
    if largest != i:
        arr[i], arr[largest] = arr[largest], arr[i]
        heapify(arr, n, largest)

@nb.njit
def push_heap(arr, val):
    arr.append(val)
    i = len(arr) - 1
    while i > 0:
        parent = (i - 1) // 2
        if arr[i][0] > arr[parent][0]:
            arr[i], arr[parent] = arr[parent], arr[i]
            i = parent
        else:
            break

@nb.njit
def pop_heap(arr):
    if not arr:
        return None
    top = arr[0]
    arr[0] = arr[-1]
    arr.pop()
    heapify(arr, len(arr), 0)
    return top

结合Numba支持的字典实现上述逻辑，可大幅提升操作速度。

方法2：使用SortedList（平衡BST实现） #

第三方库sortedcontainers中的SortedList基于平衡二叉搜索树实现，支持O(log n)时间的插入、删除和切片操作：

• 初始化：将所有(b[idx], idx)加入SortedList，直接取最后K个元素（最大值）。
• 迭代更新：每次更新b中的元素时，先删除旧的(old_b_val, idx)，再插入新的(new_b_val, idx)。
• 查找前K个最大值：直接取sorted_list[-K:]，提取其中的索引即可。

注意：若纯Python实现的SortedList在1e6次迭代下速度仍不足，可考虑用Cython或Numba实现类似的平衡BST结构，或结合分块策略优化。

问题2：其他高效方案 #

方案1：分块维护局部前K个最大值（K=1方案的扩展） #

这是最易落地且高效的方案，核心是将大数组拆分为若干小块，每个块维护自身的前K个最大值及其索引：

• 分块初始化：将b数组分成M块（比如每块1000个元素，共750块），对每个块执行argpartition(-K)[-K:]，得到该块的前K个最大值索引和对应b值，存储在块的候选列表中。
• 迭代更新：当更新某个元素时，找到其所属的块，重新计算该块的前K个最大值（或仅更新该元素在块候选列表中的位置），替换块的候选列表。
• 全局前K查找：收集所有块的候选列表（共MK个元素，比如75010=7500个），对这些候选元素执行argpartition(-K)[-K:]，快速得到全局前K个最大值索引。

该方案的优势是每次全局查找的候选集仅为几千个元素，远小于原数组的75万，argpartition速度提升数个数量级；同时更新时仅需处理单个块的候选列表，开销极低。

方案2：增量式维护全局最小堆（仅跟踪前K个最大值） #

维护一个大小为K的最小堆，堆中存储当前全局前K个最大值（堆顶是前K中的最小值）：

• 初始化：用argpartition找到初始前K个最大值，将(b[idx], idx)放入最小堆。
• 迭代更新：
  1. 若更新的元素索引在堆中：先删除旧的(old_b_val, idx)，再插入新的(new_b_val, idx)；
  2. 若更新的元素不在堆中：如果new_b_val大于堆顶的最小值，则弹出堆顶，插入新元素；
  3. 额外检查：若堆中某个元素的b值被更新后变小（不在前K了），需要从堆中移除——因K≤50，这个O(K)的开销可接受。

方案3：Numba优化全量快速选择 #

虽然全量argpartition慢，但可使用Numba自定义实现快速选择算法，仅聚焦于找出前K个最大值的索引，跳过不必要的元素遍历，相比原生numpy.argpartition可进一步减少计算开销。

内容的提问来源于stack exchange，提问作者silence_of_the_lambdas