高重复输入场景下快速选择（quickselect）算法优化方案咨询

前提说明 #

首先明确时间复杂度下限：

• 若输入为未聚合的原始数组：必须遍历全量元素完成频率统计，这一步无法突破O(n)的时间下限，空间开销固定为O(u)（仅存储唯一值与对应计数）
• 若输入为预聚合的(值, 计数)对列表：可直接跳过统计步骤，后续处理仅与唯一值数量u相关

核心优化方案：带计数的快速选择算法 #

你可以直接改造经典快速选择逻辑，适配预聚合后的唯一值计数列表，将原方案排序的O(u log u)时间复杂度降为O(u)，思路如下：

1. 先对原始数组遍历一次，用哈希表统计每个唯一值的出现次数，得到长度为u的(value, count)列表
2. 对该列表执行改造后的快速选择：
   • 随机（或用中位数的中位数规则）选取一个 pivot 值
   • 遍历列表，统计三个值：所有小于pivot的元素总计数sum_less、等于pivot的元素计数sum_eq
   • 根据k与两个累计值的关系判断递归方向：
     • 若k < sum_less：目标值在小于pivot的分组中，对该分组递归执行选择逻辑，k保持不变
     • 若sum_less ≤ k < sum_less + sum_eq：pivot就是要找的第k小值，直接返回
     • 若k ≥ sum_less + sum_eq：目标值在大于pivot的分组中，对该分组递归执行选择逻辑，k更新为k - sum_less - sum_eq

伪代码实现（Python风格） #

import random

# 输入为预聚合的(value, count)列表unique_pairs，待查询的位次k
def quickselect_with_count(unique_pairs, k):
    if len(unique_pairs) == 1:
        return unique_pairs[0][0]
    
    # 随机选pivot，也可替换为中位数的中位数实现保证最坏复杂度
    pivot_idx = random.randint(0, len(unique_pairs)-1)
    pivot = unique_pairs[pivot_idx][0]
    
    sum_less = 0
    sum_eq = 0
    less_pairs = []
    greater_pairs = []
    
    for v, c in unique_pairs:
        if v < pivot:
            sum_less += c
            less_pairs.append((v, c))
        elif v == pivot:
            sum_eq += c
        else:
            greater_pairs.append((v, c))
    
    if k < sum_less:
        return quickselect_with_count(less_pairs, k)
    elif k < sum_less + sum_eq:
        return pivot
    else:
        return quickselect_with_count(greater_pairs, k - sum_less - sum_eq)

方案优势 #

• 时间复杂度：随机选pivot时平均O(u)，替换为中位数的中位数选pivot规则可保证最坏O(u)，远优于原方案1的O(u log u)排序开销
• 空间复杂度：O(u)，仅存储唯一值与对应计数，没有方案2的O(M)前置开销，不需要提前知晓数值上下限，也不要求数值为整数类型
• 适配你的业务场景：u基本恒定的前提下，聚合后的查询耗时不会随n的波动变化，性能稳定性极高

预聚合输入场景适配 #

如果输入已经是完成聚合的(v1, c1), (v2, c2)...格式，可直接跳过哈希表统计步骤，直接调用上述快速选择逻辑即可，总时间复杂度直接为O(u)，无需遍历原始全量n个元素。

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