使用Python打算法竞赛的小技巧（持续更新）

使用 Python 打算法竞赛似乎一直以来不被看好，原因主要是 Python 本身的运行效率不如 C++（尤其是只提供 CPython 时），且容易在一些好用的 API 或包的诱惑下写出大常数代码，在算法竞赛这个寸时寸金的领域稍有不慎就会 TLE。有些题可能 C++ 放过了 $O(n\log n)$ 的做法，但 Python 必须写 $O(n)$ 做法才能过。随着时代的发展，编程的普及，越来越多的朋友使用 Python 写算法题，然后就不得不直面 Python 慢的问题。在这个过程中，大家探索出了很多可以让 Python 跑得更快的实现技巧，Python 逐渐地在一些题型上有了一战之力，在蓝桥杯等比赛中甚至也出现了只能使用 Python 的组别。

本文主要介绍 Python 在做算法题时如何减小 Python 代码的常数（因为讲基本用法的文章网上已经有很多了），从而有可能通过更多的题目。假如比赛环境中支持多种语言，但 Python 并没有提供 PyPy 的话，最好还是别用 Python 了，老老实实用 C++ 吧。

输入输出

输入

一般情况下，使用 input() 函数就可以了，常见的输入模式：

# 单个数读入
n = int(input())
# 同行读入若干个数
n, k = map(int, input().split())
# 同行读入一个数组
a = list(map(int, input().split()))
# 同行读入一个数组，并且强行以 1 作为起始下标（如果你更习惯 1 作为起始下标）
a = [0] + list(map(int, input().split()))

但是，在需要大量输入数据的题目中，input 会处理内置的输入缓冲（包括提示符、换行处理等），直接使用 input() 函数就会 TLE 了，这个时候我们需要加这样一行代码：

input = lambda: sys.stdin.readline().strip()

加上它之后，后边的读入代码完全不用变。

sys.stdin.readline() 更快，因为它直接使用系统级输入读取，但是保留了换行符，需要 strip() 去掉。

输出

主要是使用 print() 函数，但是每调用一次 print() 函数就是一次 I/O 操作，I/O 操作过多会导致速度变慢。

我们考虑以下代码：

n = int(4e6)

for i in range(1, n + 1):
    print(i)

在 Codeforces custom invocation 中，使用 PyPy 3.10 进行测试，其运行了 687ms，多次运行均差不多是这些。

如果我们把所有的输出塞到列表里，最后统一输出：

n = int(4e6)
res = []
for i in range(1, n + 1):
    res.append(i)
print('\n'.join(map(str, res)))

需要 468ms，多次测试均差不多是这些。

上面只是一个简单的例子。事实上，在其他的需要大量输出的场景下，把输出最后统一进行总是会快一些。因此，在内存允许的情况下，可以把数据结构题或者多测题的输出攒起来，最后调用一次 print() 函数一起输出，能够提升一些速度。

排序

对于一般的整数数组，我们可以直接使用 nums.sort() 进行排序。

有时候我们需要对一些复杂对象的列表进行排序，这个时候我们通过定义比较函数，使用 cmp_to_key 将比较函数转化为 key 之后进行排序即可。

从实践上来说，通过指定 key 直接对复杂对象列表排序会带来很多拷贝的开销，导致速度变慢很多。

这时，我们考虑另一个思路：我们初始化一个下标列表 [1, n]，然后对下标列表进行排序，排序的 key 中使用了复杂对象列表中的数据。这样之后，我们会得到一个下标列表，按照这个列表从左往右遍历，这些下标对应的复杂对象列表中的数据就是排序后的顺序。

看一份代码理解一下吧。这里是把元组 (x, y) 按照 x 从小到大排序，y 的顺序不关心，使用对下标排序的方案实现：

m = int(input())
xs = []
ys = []
for i in range(m):
    x, y = map(int, input().split())
    xs.append(x)
    ys.append(y)
 
order = [i for i in range(m)]
order = sorted(order, key=lambda i: xs[i])

由于不涉及原数据的修改，所以我们甚至不需要真的使用元组去存，而是可以直接俩 list 对齐去存储。

对下标进行排序的方案我们只是对一个整数数组进行的排序，得到的相当于复杂对象列表的一个排序视图，除了快之外，我们还可以在原列表的基础上按照不同的排序规则导出多个排序视图。

关于排序的一道例题可以看 CF652C。在本题中，你可以清楚地感知到快速读入、复杂对象排序、哈希等对 Python 运行速度的影响。

多维数组

Python 实现类似多维数组的东西，最直观的写法是 list 套 list。但这个东西非常慢， $n=5000$ 有时候都会寄。list 第一维远大于第二维的时候（例如某些 ST 表或倍增的写法），很容易出现类似 C++ 的 Cache miss 的现象，且慢得更加明显。

在 DP 时，能用滚动数组就用滚动数组，这样少一维不只是空间优化了，时间也能快一些。

list 套 list 还可以考虑拍平成一维的 list（这正好是 C 语言编译时对高维数组的处理方式），可以写一个下标变换函数 f，把 dp[i][j][k] 这种高维下标变成一维下标 dp[f(i, j, k)]。

拍平下标转化也是有时间开销的。如果是 DP 场景的话，可能只涉及加法和乘法，开销还好。但有的题需要用除法和取模去把一维下标还原成多维下标，这个时候就不太好说了。

手写工具函数

在 LeetCode 周赛中，经常有朋友发现因为使用了 max 之类的函数而超时，这个时候需要改成手写 fmax 函数才能通过：

fmax = lambda a, b: a if a > b else b

这个坑点在 LeetCode 上最明显，说是能差出 3 倍；在 Codeforces 上我测试时感觉差别比较小。

值得一提的是，我在 Codeforces 上测试 max 时，对于相同的代码，选择 Python 3.10 和选择 PyPy 3 运行的时间分别是 600+ms 和 93ms，可见 PyPy 3 在算法竞赛中对 Python 有多重要。

有时候 max 只是让程序 TLE 的最后一根稻草，或已经 TLE 之后的雪上加霜，这时候还需要看其他地方有没有可以优化的点。

非递归

树相关问题的非递归

直接使用 DFS 进行树的遍历，很有可能会爆栈 RE，改了栈空间大小可能也不好使（TLE MLE RE 快乐三选一）。有人说我们可以写 BFS 遍历呀，但是有时候我们就是需要按照自底向上的那个顺序去遍历结点（比如需要利用子树信息更新本结点信息的树形 DP），所以我们还是要想办法知道 DFS 的结点遍历顺序。

注意到 BFS 是一层一层访问的，其访问顺序反过来的话，是一种合理的自底向上的顺序，所以我们使用 BFS 去遍历树，得到一个序列 order，并记录 BFS 时每个结点的父亲，然后按照 order 的反序去遍历，就可以自底向上去做一些事情（子结点信息更新父亲信息）；按照 order 的正序遍历，则可以自顶向下做一些事情（父结点信息更新子结点信息）。

代码如下：

q = []
order = []
q.append(1)
parent = [-1 for _ in range(n + 1)]
parent[1] = 0
front = 0
while front < len(q):
    u = q[front]
    front += 1
    order.append(u)
    for v in e[u]:
        if parent[v] == -1:
            parent[v] = u
            q.append(v)

reversed_order = order[::-1]

for u in reversed_order:
    ...

事实上，可以直接用 order 作为 BFS 的队列，无需单独开一个 q。

掌握了该项技能，就又可以过很多题了。

数据结构

对 Python 速度的正确感知

作用域相关的坑

参考文献

灵茶山艾府、小羊肖恩、conqueror_of_tourist 等经常使用 Python 做题/讲题的大佬的 AC 代码。
算法交流群里关于 Python 使用的聊天记录。
本人的实践经验。