《Python算法竞赛攻略：从入门到降维打击》——第十二章

第12章：Pythonic效率工具箱

12.1 itertools: 排列、组合、笛卡尔积的数学之美

在算法竞赛中，我们经常遇到涉及排列组合的搜索问题。例如，“从N个物品中选取M个的所有方案”、“对一组元素进行全排列”等等。在C++中，这通常需要我们手写复杂的递归函数，不仅耗时，还容易出错。而Python的itertools模块，将这些基础的数学概念封装成了高效、简洁的工具，让我们能一行代码解决问题，专注于算法逻辑本身。

核心理念： 不要手动实现轮子，itertools就是你最可靠的数学工具。

1. 笛卡尔积 (itertools.product)

功能： 计算多个可迭代对象的笛卡尔积。它等价于多重嵌套循环。

应用场景： 当你需要从多个独立的集合中各取一个元素，组成所有可能的组合时。例如，多颗骰子的点数组合、多件不同部位装备的搭配方案等。

实战案例：
假设有3个密码盘，第一个盘有数字[1, 2]，第二个盘有字母['A', 'B']，第三个盘有符号['#', '*']。求所有可能的密码组合。

import itertools

p1 = [1, 2]
p2 = ['A', 'B']
p3 = ['#', '*']

# 使用itertools.product计算笛卡尔积
for password in itertools.product(p1, p2, p3):
    print(password)

# 输出:
# (1, 'A', '#')
# (1, 'A', '*')
# (1, 'B', '#')
# (1, 'B', '*')
# (2, 'A', '#')
# (2, 'A', '*')
# (2, 'B', '#')
# (2, 'B', '*')

降维打击效果： 仅用一行itertools.product(p1, p2, p3)就替代了三层嵌套的for循环，代码意图一目了然。

2. 排列 (itertools.permutations)

功能： 从一个可迭代对象中，取出指定数量的元素进行全排列。

应用场景： 解决需要考虑“顺序”的问题。例如，N个人排队的所有排法、用数字1,2,3组成所有可能的三位数。

实战案例：
给定一个数组[1, 2, 3]，输出其所有长度为3的全排列。

import itertools

nums = [1, 2, 3]

# 计算长度为3的全排列
for p in itertools.permutations(nums, 3): # r参数可选，默认为len(nums)
    print(p)

# 输出:
# (1, 2, 3)
# (1, 3, 2)
# (2, 1, 3)
# (2, 3, 1)
# (3, 1, 2)
# (3, 2, 1)

降维打击效果： 无需编写递归和visited数组，一行代码即可优雅地生成全排列。

3. 组合 (itertools.combinations)

功能： 从一个可迭代对象中，取出指定数量的元素进行组合，不考虑顺序。

应用场景： 解决不考虑“顺序”的选取问题。例如，从N个人中选M个人组成委员会、从一堆数中选出两个数求和。

实战案例：
在一个数组[1, 2, 3, 4]中，找出所有大小为2的组合。

import itertools

nums = [1, 2, 3, 4]

# 计算大小为2的组合
for c in itertools.combinations(nums, 2):
    print(c)

# 输出:
# (1, 2)
# (1, 3)
# (1, 4)
# (2, 3)
# (2, 4)
# (3, 4)

降维打击效果： 对于组合问题，itertools.combinations是无可替代的最优解，避免了为保证顺序而设计的复杂循环或递归。

4. 可重复组合 (itertools.combinations_with_replacement)

功能： 类似组合，但允许同一个元素被多次选择。
应用场景： 可重复的选取问题。例如，从3种口味的冰淇淋中买2个球，可以买两个相同口味的。

竞赛提示：
itertools返回的是一个迭代器，而不是列表。这意味着它在内存使用上极为高效，只有在你遍历它的时候才会逐个生成元素。但是，如果排列组合的总数巨大（例如对20个元素全排列），即使使用itertools，程序也必然会超时。在使用前，请务必先用数学公式（如AnmA_n^mAnm​, CnmC_n^mCnm​）估算结果集的规模，判断是否会在时限内完成。

---

12.2 heapq: 优先队列、Top-K问题、堆排序的利器

heapq模块是Python对堆这种数据结构的实现。它不是一个类，而是一系列作用于列表上的函数。heapq实现的是最小堆，即堆顶元素heap[0]永远是最小的。这使得它成为实现优先队列、解决Top-K问题和优化某些图算法的绝佳工具。

核心理念： 当你需要动态维护一组数据中的最小值或最大值时，第一时间想到heapq。

1. 核心操作

heapq.heappush(heap, item): 将item压入堆heap中，并保持堆的性质。时间复杂度为 O(log⁡N)O(log N)O(logN)。
heapq.heappop(heap): 从堆heap中弹出并返回最小的元素。时间复杂度为 O(log⁡N)O(log N)O(logN)。
heapq.heapify(x): 将列表x原地转换为一个最小堆。时间复杂度为 O(N)O(N)O(N)。

2. 实战应用：Top-K 问题

问题描述： 从一个巨大的数据流或数组中，找出前K个最大（或最小）的元素。

降维打击思路：

求Top-K最大值： 维护一个大小为K的最小堆。遍历所有元素，如果当前元素比堆顶（堆中最小的元素）还大，则弹出堆顶，将当前元素压入。遍历结束后，堆中剩下的K个元素就是最大的K个。
求Top-K最小值： 同理，维护一个大小为K的最大堆。

如何用最小堆模拟最大堆？
一个非常Pythonic的技巧：存入元素的相反数。这样，原本最大的元素就变成了最小的负数，会被放在最小堆的堆顶。

实战案例：
找出数组[3, 5, 1, 7, 6, 8, 2, 9, 4, 0]中最大的3个元素。

import heapq

nums = [3, 5, 1, 7, 6, 8, 2, 9, 4, 0]
k = 3
min_heap = []

for num in nums:
    if len(min_heap) < k:
        heapq.heappush(min_heap, num)
    elif num > min_heap[0]: # 当前元素比堆顶（已有的k个最大数中的最小那个）还大
        heapq.heappop(min_heap) # 弹出最小的
        heapq.heappush(min_heap, num) # 压入更大的

# heappop(min_heap)之后再heappush(min_heap, num)可以被一个函数替代
# heapq.heapreplace(min_heap, num)

print(f"最大的 {
                k} 个数是: {
                sorted(min_heap, reverse=True)}")
# 输出: 最大的 3 个数是: [9, 8, 7]

3. 在图论算法中的应用

正如第9章所提到的，heapq是Dijkstra算法和Prim算法的性能关键。通过将(距离, 顶点)这样的元组存入堆中，heapq可以自动根据距离（元组的第一个元素）进行排序，让我们在每一步都能以 O(log⁡N)O(log N)O(logN) 的代价快速获取当前距离最小的顶点。

4. 便捷函数

heapq.nlargest(k, iterable) / heapq.nsmallest(k, iterable): 如果你只是需要一次性地从一个静态集合中找出Top-K个元素，这两个函数是更直接的选择，其内部实现正是我们上面讨论的堆方法。

---

12.3 functools与operator: @cache再探与高阶函数利器

functools 和 operator 是Python函数式编程思想的体现。它们提供了许多高阶函数和可调用对象，能让我们的代码更加声明式和高效。

1. @cache 再探

我们在第8章动态规划中已经见识过@functools.cache的威力。它是一键开启“记忆化搜索”的神器，能自动缓存函数对于特定参数的返回值。这里我们再次强调它的重要性：任何时候当你写出一个纯函数（相同输入总有相同输出）的递归，并且发现有大量重复计算时，请毫不犹豫地给它戴上@cache的帽子。这是Python在解决DP问题时最显著的优势之一。

2. operator.itemgetter: 更快更优雅的排序键

在排序（第5章）时，我们常用lambda函数来指定排序的关键字。

data = [('Alice', 25), ('Bob', 20), ('Cidy', 25)]
# 按年龄排序，年龄相同按名字排序
data.sort(key=lambda x: (x[1], x[0])) 

operator.itemgetter可以实现同样的功能，并且通常比lambda更快，代码也更具可读性。

from operator import itemgetter

data = [('Alice', 25), ('Bob', 20), ('Cidy', 25)]
# itemgetter(1) 表示获取索引为1的元素（年龄）
# itemgetter(1, 0) 表示先按索引1，再按索引0
data.sort(key=itemgetter(1, 0)) 

itemgetter(1, 0)会创建一个可调用对象，其功能等价于lambda x: (x[1], x[0])。对于需要访问对象属性的场景，还有operator.attrgetter。

3. functools.cmp_to_key: 兼容C++的cmp式排序

Python 3取消了sort方法中的cmp参数，统一使用key。但在某些特殊的排序逻辑中，定义两个元素a和b的直接比较关系cmp(a, b)，比定义单个元素的映射值key(a)要自然得多。

应用场景： LeetCode经典题“最大数”。给定一组非负整数，重新排列它们，使其组成的整数最大。例如，[3, 30, 34, 5, 9]应排列为9534330。

排序规则： 对于任意两个数（字符串形式）a和b，如果a+b > b+a，则a应该排在b的前面。

cmp_to_key的用法：

定义一个老式的cmp函数，它接收两个参数a, b，如果a < b返回负数，a == b返回0，a > b返回正数。
使用functools.cmp_to_key将这个cmp函数转换成一个可用的key函数。

实战案例：

import functools

def largestNumber(nums) -> str:
    # 定义比较函数
    def compare(a, b):
        if a + b > b + a:
            return -1 # a排在前面
        elif a + b < b + a:
            return 1  # b排在前面
        else:
            return 0

    # 将数字转为字符串
    s_nums = [str(n) for n in nums]

    # 使用cmp_to_key进行排序
    s_nums.sort(key=functools.cmp_to_key(compare))

    # 处理全为0的特殊情况
    if s_nums[0] == '0':
        return '0'
    
    return "".join(s_nums)

print(largestNumber([3, 30, 34, 5, 9]))
# 输出: 9534330

降维打击效果： cmp_to_key为你打开了一扇门，让你能用最直观的方式定义复杂的元素间排序关系，而不必绞尽脑汁去构造一个等价的key函数。

---

12.4 Python原生大整数：告别高精度计算的烦恼

这或许是Python给C++选手带来的最大“降维打击”，简单到甚至不需要引入任何模块。在C++或Java中，一旦数值超出了long long或long的范围，你就必须手写高精度算法，或者引入第三方大数库，过程繁琐且易错。

在Python中，整数类型可以表示任意大小的整数，仅受限于你的机器内存。

核心理念： 当题目涉及大数运算时，直接算，Python会帮你处理一切。

1. 阶乘与幂运算

计算100! (100的阶乘) 或者 2**1024，在C++中是典型的高精度题目，在Python中：

import math

# 计算100的阶乘
large_factorial = math.factorial(100)
print(f"100! = {
              large_factorial}")

# 计算2的1024次方
large_power = 2**1024
print(f"2**1024 = {
              large_power}")

代码会直接输出这两个巨大的数字，你不需要为它们声明任何特殊类型。

2. 组合数学中的应用

在计算组合数 Cnk=n!k!(n−k)!C_n^k = frac{n!}{k!(n-k)!}Cnk​=k!(n−k)!n!​ 时，中间过程的阶乘n!可能非常大，但最终结果可能在64位整数范围内。在C++中，你需要边乘边除来避免溢出。在Python中，你可以直接计算，完全不用担心中间过程。

import math

# 使用math.comb计算组合数 C(100, 50)
# 这个函数内部也受益于大整数支持
print(math.comb(100, 50))

# 手动计算，同样不会溢出
n, k = 100, 50
result = math.factorial(n) // (math.factorial(k) * math.factorial(n - k))
print(result)

3. 模运算中的妙用

我们在第11章介绍过的pow(a, b, m)函数，是快速幂模运算的终极利器。这里的a和b同样可以是Python的大整数。这意味着你可以计算 (123456789987654321)(mod109+7)(123456789^{987654321}) pmod{10^9+7}(123456789987654321)(mod109+7) 这样的问题，而无需担心底数和指数本身的大小。Python的数学运算内核会高效地处理这一切。

总结

本章介绍的itertools, heapq, functools, operator以及原生大整数，是Python标准库中的瑰宝。它们是“Pythonic”思想在算法竞赛中的具体体现。掌握它们，意味着你不仅拥有了解决问题的算法知识，更拥有了选择最佳工具的智慧。这正是从“会写代码”到“写好代码”，从“暴力AC”到“优雅降维打击”的关键一步。在后续的实战中，请时刻提醒自己：当遇到相应场景时，我是否想起了这些Pythonic的效率工具？