Python求解Kattis knigsoftheforest问题超时，求优化方案及竞赛语言建议

问题原因分析 #

你两次写的代码超时的核心原因都是时间复杂度达到了O(n²)，完全无法适配n和k最大1e5的规模：

• 第一版代码：每年遍历全量列表筛选符合参赛年份的驼鹿，加上列表remove操作本身也是O(n)，n次循环总复杂度直接到平方级
• 第二版代码：每年遍历全量字典筛选手法，还多次复制字典，本质还是暴力遍历，哪怕用了堆也没有降低整体复杂度

优化方案 #

正确的思路可以把时间复杂度降到O((n+k)log(n+k))，完全满足1秒时间限制：

1. 先把所有驼鹿按参赛年份从小到大排序，用一个指针跟踪还没加入参赛池的驼鹿，避免每次遍历全量数据
2. 用最大堆（Python的heapq是最小堆，所以存力量的负值模拟最大堆）维护当前参赛池的所有驼鹿力量值
3. 逐年模拟比赛：
   • 一次性把当年新参赛的所有驼鹿推入堆
   • 弹出堆顶的最大力量值作为冠军，如果是Karl直接输出当前年份
   • 模拟完所有已知年份还没轮到Karl获胜就输出unknown
4. 额外做输入优化：Python原生input读1e5次数据非常慢，直接用sys.stdin.read()一次性读全量数据再拆分，能大幅提升输入速度

优化后代码 #

import sys
import heapq

def main():
    data = list(map(int, sys.stdin.read().split()))
    ptr = 0
    k, n = data[ptr], data[ptr+1]
    ptr += 2
    karl_y, karl_p = data[ptr], data[ptr+1]
    ptr += 2
    
    moose_list = []
    moose_list.append((karl_y, karl_p))
    for _ in range(n + k - 2):
        y, p = data[ptr], data[ptr+1]
        moose_list.append((y, p))
        ptr += 2
    
    # 按参赛年份排序所有驼鹿
    moose_list.sort()
    heap = []
    moose_ptr = 0 # 标记还没进参赛池的驼鹿下标
    
    for year in range(2011, 2011 + n):
        # 把当年所有新参赛的驼鹿全部加入堆
        while moose_ptr < len(moose_list) and moose_list[moose_ptr][0] == year:
            heapq.heappush(heap, -moose_list[moose_ptr][1])
            moose_ptr += 1
        # 选出当年冠军
        winner_p = -heapq.heappop(heap)
        if winner_p == karl_p:
            print(year)
            return
    print("unknown")

if __name__ == "__main__":
    main()

编程语言选择建议 #

• 如果你的目标是冲击USACO白金组或者更高级别的算法竞赛，非常推荐学习C++：同等算法下C++运行速度是Python的10~20倍，卡时间阈值的题不容易超时，而且是竞赛的主流语言，学习资源、样例代码都更丰富。
• 如果只是打铜/银组，或者更习惯Python的语法，只要注意选择正确复杂度的算法、做好输入输出优化，大部分题也都能通过。

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