Python求解Grandpa Bernie旅行查询问题的速度优化咨询

题目说明 #

多年来，Bernie爷爷环游了世界各地，如今他不再频繁出行，但很喜欢给孙辈讲述过往旅行的故事，比如他第一次去以色列、第三次去希腊的经历。他的记忆模式很特别：能轻松记起到某一国家的第k次旅行，却很难记起那次出行的年份。给定Bernie爷爷所有旅行的列表，你需要处理多个查询，返回他到指定国家的第k次旅行的年份。

输入规则 #

• 第一行是一个整数n（1≤n≤100000），代表Bernie爷爷的旅行总次数；
• 接下来n行，每行包含国家名s（长度1≤|s|≤20，仅由英文字母组成）和整数y（1≤y≤1000000），代表一次去s国的旅行发生在y年；
• 接下来一行是一个整数q（1≤q≤100000），代表查询总次数；
• 接下来q行，每行包含国家名s和整数k，代表查询去s国的第k次旅行的年份。题目保证所有查询的k合法，取值范围为1到s国的旅行总次数，且s国至少被访问过一次。

输出规则 #

每个查询单独输出一行，对应到指定国家第k次旅行的年份。

原有代码问题分析 #

你的解题思路是正确的，但是实现存在两个核心的性能瓶颈：

1. 查询时重复排序：每次查询都对对应国家的年份列表执行一次sorted()操作，当查询量达到1e5级别时，会产生大量重复计算，时间复杂度直接飙升到不可接受的范围。
2. IO效率过低：Python原生的input()和print()函数在处理十万级别的输入输出时，因为每次调用都有系统开销，速度非常慢。

优化方案 #

优化1：预处理阶段统一排序 #

所有旅行记录读取完成后，只需要遍历一次字典，给每个国家的年份列表做一次升序排序，后续查询直接取下标即可，排序总时间复杂度为O(n log n)，远低于查询时重复排序的开销。

优化2：替换高开销的IO操作 #

使用sys.stdin.readline()替代input()读取输入，输出时先把所有结果存入列表，最后一次性拼接打印，减少IO调用次数。

优化3：使用defaultdict简化字典操作 #

使用collections.defaultdict可以省去判断国家是否存在于字典中的逻辑，稍微提升代码运行效率。

优化后代码 #

import sys
from collections import defaultdict

def main():
    input = sys.stdin.readline
    dates = defaultdict(list)
    n = int(input())
    for _ in range(n):
        s, y = input().split()
        dates[s].append(int(y))
    # 预处理：每个国家的年份列表只排序一次
    for s in dates:
        dates[s].sort()
    q = int(input())
    res = []
    for _ in range(q):
        s, k = input().split()
        res.append(str(dates[s][int(k)-1]))
    # 一次性输出所有结果
    print('\n'.join(res))

if __name__ == "__main__":
    main()

优化效果说明 #

• 排序开销从O(q * k log k)（k为单个国家旅行次数）降低到O(n log n)，十万级数据下排序耗时减少99%以上。
• IO操作次数从2*(n+q)次降低到3次左右，大幅降低系统调用开销。
  优化后的代码可以轻松在1秒时间限制内跑完最大规模的测试用例。

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