Python数据结构问题求解：统计文件与Top K集合是否需要用heap？

核心思路拆解 #

我们可以把需求拆成两个独立的逻辑模块处理，整体时间复杂度最优可到O(N + M log K)，其中N为文件总数，M为不同集合总数，K为指定的Top排名数量：

• 模块1：总文件大小统计，只需要维护一个全局累加变量即可，每处理一个文件直接累加文件大小，不受集合归属影响。
• 模块2：集合总大小统计+Top K计算，分两步走：
  1. 先用哈希表（字典）统计每个集合的累计总大小：遍历所有文件，若文件有归属集合，就把文件大小累加到每个所属集合的计数上，天然支持单文件归属多集合的规则。
  2. 再用小根堆（最小优先队列）实现Top K筛选，你判断的方向完全正确，这是该场景下的最优方案之一。

小根堆的具体实现逻辑 #

优先选择小根堆而非大根堆的原因是可以用更低的空间开销完成筛选：维护一个固定大小为K的小根堆，堆顶始终存储当前已遍历集合中总大小最小的那个，遍历所有集合的总大小的时候执行以下判断：

• 如果堆的当前长度小于K，直接将（集合总大小，集合名）推入堆中
• 如果堆长度等于K，且当前集合的总大小大于堆顶的大小，弹出堆顶元素，将当前集合信息推入堆中
  遍历完成后，堆内存储的就是总大小排名前K的集合，最后按大小倒序输出即可。

边界情况处理：如果集合总数小于K，直接输出所有集合即可；如果没有有效集合，Top K部分可以留空。

参考代码示例（Python） #

import heapq
from collections import defaultdict

def solve(file_list, K):
    total_size = 0
    collection_size = defaultdict(int)
    
    # 第一步：遍历所有文件统计总大小和集合累计大小
    for file_name, file_size, collections in file_list:
        total_size += file_size
        for coll in collections:
            collection_size[coll] += file_size
    
    # 第二步：小根堆求Top K
    heap = []
    for coll_name, size in collection_size.items():
        if len(heap) < K:
            heapq.heappush(heap, (size, coll_name))
        else:
            if size > heap[0][0]:
                heapq.heappop(heap)
                heapq.heappush(heap, (size, coll_name))
    
    # 堆内默认从小到大排序，输出要求从大到小所以反转结果
    top_k = sorted(heap, reverse=True)
    
    # 格式化输出
    print(f"Total size of files processed: {total_size}")
    print(f"Top {K} collections:")
    for size, name in top_k:
        print(f"-{name}: {size}")

# 代入示例输入测试
if __name__ == "__main__":
    K = 2
    # 输入按[文件名, 文件大小, 归属集合列表]格式化
    file_list = [
        ["file1.txt", 100, []],
        ["file2.txt", 200, ["collection1"]],
        ["file3.txt", 200, ["collection1"]],
        ["file4.txt", 300, ["collection2"]],
        ["file5.txt", 100, []],
        ["file6.txt", 200, ["collection1", "collection3"]]
    ]
    solve(file_list, K)

运行上述代码即可得到和示例完全一致的输出。

复杂度对比 #

如果不用堆，直接对所有集合排序取Top K，时间复杂度为O(M log M)，当M非常大（比如百万级以上）、K很小（比如Top 10）的时候，堆方案的O(M log K)效率要高几十上百倍，性能优势非常明显。

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