如何对多组数值数组进行计数排序并在O(n)时间内分别返回各组结果

现有m个存储了1~n范围内数值的数组，所有数组的元素总个数为n，即size(m1) + size(m2) + … + size(mm) = n，需要对每个数组分别排序后单独返回。要求算法时间复杂度为O(n)，不可为O(m+n)。常规的合并所有数组后执行计数排序的方案虽然能达到O(n)时间复杂度，但只能得到所有元素合并排序后的单个数组，无法分组返回每个原数组的排序结果。

直接改造计数排序的逻辑即可，不需要做数组合并，全程时间复杂度严格为O(n)。

操作步骤 #

• 给每个原数组分配一个唯一的序号，比如从0到m-1。
• 初始化长度为n+1的桶数组（适配1~n的数值范围），每个桶是一个空列表，用来存对应数值所属的原数组序号。
• 遍历所有原数组的所有元素：如果当前遍历的是第k个数组的元素x，就把k追加到桶x的末尾，这一步总遍历次数是n，时间复杂度O(n)。
• 给每个原数组初始化一个写指针，初始值为0，用来标记当前排序结果的写入位置。
• 按从小到大的顺序遍历1到n的所有数值x：遍历桶x里存的所有原数组序号k，把x写到第k个数组结果的写指针位置，然后把第k个数组的写指针加1。这一步遍历数值的次数是n，所有桶的元素总数也是n，总操作次数还是O(n)。

遍历完成后，每个原数组对应的结果就是单独排序好的数组。

伪代码示例 #

# 输入：arrays = [arr_0, arr_1, ..., arr_{m-1}]，数值范围1~n，总元素数为n
# 1. 初始化桶
buckets = [[] for _ in range(n + 1)]
# 2. 统计每个数值所属的原数组
for k in range(len(arrays)):
    for x in arrays[k]:
        buckets[x].append(k)
# 3. 按数值从小到大回写到对应数组
res = [[0] * len(arr) for arr in arrays]
ptr = [0] * len(arrays)
for x in range(1, n + 1):
    for k in buckets[x]:
        res[k][ptr[k]] = x
        ptr[k] += 1
# res就是最终结果，每个子数组对应原数组的排序结果

复杂度说明 #

整个算法的所有遍历操作总次数严格等于2n，就算考虑初始化结果数组、指针数组的开销，由于m最大不超过n（极端情况所有数组都只有1个元素），这部分开销也属于O(n)范畴，整体复杂度严格为O(n)，完全符合要求。

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