查找最近未取用元素的数据结构及线性复杂度实现方案

给定长度为N的数组A[0..N-1]，以及长度为M的查询索引列表，需要按从前往后的顺序处理每个查询索引i，取用规则如下：

• 若A[i]未被取用，直接取用A[i]，本次取用的索引记为i
• 若A[i]已被取用，查找满足j > i的最小未被取用索引j，取用A[j]，本次取用的索引记为j；如果不存在符合要求的j，本次结果记为-1
  最终输出长度为M的结果数组B，B[k]对应第k次处理时取用的元素索引，要求整体实现达到线性时间复杂度（O(N)或O(M)级别），需要选择适配的数据结构。

这个问题要做到线性级别的复杂度，选带路径压缩的并查集就够了，不用上线段树、平衡树这类带log复杂度的结构，实现简单运行效率也更高。

核心逻辑 #

本质上我们要快速找到「大于等于查询值i的最小未被占用索引」，用并查集维护每个位置对应的下一个可用位置即可：

• 初始化的时候，每个位置的父节点指向自己，额外多设一个索引为N的哨兵节点（数组最大有效索引是N-1），哨兵的父节点指向自己，代表找到这里就没有可用元素了
• 处理查询i的时候，顺着父节点找根节点，这个根节点就是i后面第一个可用的索引
• 找到可用索引root之后，直接把root的父节点设为root+1的根节点，相当于标记root已经被占用，下次再找的时候直接跳过它
• 查找过程中做路径压缩，把沿途经过的所有节点的父节点直接指向最终找到的根，后续查询不用重复遍历路径

参考实现 #

def process_queries(N: int, query_list: list[int]) -> list[int]:
    # parent数组长度为N+1，索引N作为无可用元素的哨兵
    parent = list(range(N + 1))
    
    def find(x: int) -> int:
        if parent[x] != x:
            parent[x] = find(parent[x]) # 路径压缩
        return parent[x]
    
    res = []
    for i in query_list:
        available_idx = find(i)
        if available_idx == N:
            res.append(-1)
        else:
            res.append(available_idx)
            # 标记当前索引已占用，指向下一个位置的根
            parent[available_idx] = find(available_idx + 1)
    return res

复杂度说明 #

带路径压缩的并查集单次操作的均摊时间复杂度是反阿克曼函数级别，实际运行中可以认为是常数操作，整体总时间复杂度为O(N+M)，完全满足线性复杂度要求，额外空间复杂度为O(N)，同样是线性级别。

小提示：这个实现里的哨兵节点设计省掉了很多边界判断逻辑，不用单独处理i是最后一个位置、后面没有可用元素的场景，代码更简洁不容易出bug。

内容的提问来源于stack exchange，提问作者Y.T.