Python tricks

Python 越来越多地成为大家刷题的主流语言，主要原因是它的语法非常简洁明了。因此我们能节省更多的时间，来关注算法和数据结构本身。

而用好 Python 自身独有的一些语法特性，不仅能更节省时间，也能让代码看起来更加优雅。这里我总结了一些我自己刷题过程中用到的一些常用的功能。以下均以 Python 3 为例（Python 2 已于 2020 年停止维护，不再建议使用）。

List

Python 的列表 List 基本就是其它语言的 Array.

Initialization 初始化

List 的初始化一般用 List comprehension，往往能一行解决问题：

# 1d array（长度为 n，全部初始化为 0）
l = [0] * n
# 或者
l = [0 for _ in range(n)]

# 2d array（rows 行 cols 列，全部初始化为 0）
l = [[0] * cols for _ in range(rows)]

注意 2d 数组不要写成 [[0] * cols] * rows，因为每一行会指向同一个 list 对象，修改一行会影响到所有行。

索引与切片

你可以轻松从后往前访问，也支持步长：

lastElement = l[-1]      # 最后一个元素
lastTwo = l[-2:]         # 最后两个元素
firstTwo = l[:2]         # 前两个元素
l[::-1]                  # 反转整个列表
l[::2]                   # 每隔一个元素取一个

for i in range(0, -10, -1):
    print(i)
# 0, -1, -2, -3, -4, -5, -6, -7, -8, -9

解包 unpacking

* 和 ** 可以方便地展开可迭代对象和字典：

first, *rest = [1, 2, 3, 4]
# first = 1, rest = [2, 3, 4]

a, *mid, b = [1, 2, 3, 4, 5]
# a = 1, b = 5, mid = [2, 3, 4]

merged = [*l1, *l2]      # 合并两个列表
d = {**d1, **d2}         # 合并两个字典

copy 复制

shallow copy 浅拷贝

l2 = l1[:]
# or
l2 = l1.copy()

浅复制的问题在于，如果 l1 内部还有 list，那么这种嵌套的索引不能被复制，比如：

a = [1, 2, [3, 4]]
b = a[:]
a[2].append(5)
print(b)
# [1, 2, [3, 4, 5]]

deep copy 深拷贝

如果要做深拷贝，要借助自带库 copy：

import copy

b = copy.deepcopy(a)

enumerate 枚举

当我们需要枚举一个数组并同时获得值与 index 的时候可以使用，start 参数可以指定起始下标：

l = ["a", "b", "c"]

for i, v in enumerate(l):
    print(i, v)
# 0 a
# 1 b
# 2 c

for i, v in enumerate(l, start=1):  # 从 1 开始
    print(i, v)

zip

zip 本意就是拉链，可以想象成将两个数组像拉链一样挨个聚合：

>>> x = [1, 2, 3]
>>> y = [4, 5, 6]
>>> zipped = zip(x, y)
>>> list(zipped)
[(1, 4), (2, 5), (3, 6)]

• 用 zip(*matrix) 可以方便地得到矩阵的转置。
• 如果两个列表长度不一致，默认会截断到较短的；若希望长度不一致时直接报错，可传 strict=True（Python 3.10+）。
• 需要补长到最长时，用 itertools.zip_longest。

reduce

reduce 可以分别对相邻元素使用同一种计算规则，同时每一步结果作为下一步的参数，很典型的函数式编程用法。

from functools import reduce

lis = [1, 3, 5, 6, 2]

print(reduce(lambda a, b: a + b, lis))
# 17

map / filter

map 可以将参数一一映射来计算，比如：

date = "2019-8-15"
Y, M, D = map(int, date.split('-'))
# Y = 2019, M = 8, D = 15

nums = list(map(str, [1, 2, 3]))   # ['1', '2', '3']

filter 则按条件过滤元素：

even = list(filter(lambda x: x % 2 == 0, [1, 2, 3, 4, 5]))
# [2, 4]

deque

list 删除末尾的操作是O(1)的，但是删除头操作就是O(n)，这时候我们就需要一个双端队列 deque。首尾的常规操作为：

• append，添加到末尾
• appendleft, 添加到开头
• pop, 剔除末尾
• popleft，移除开头

sorted

list 自身有自带的 sort(), 但是它不返回新的 list. sorted 能返回一个新的 list, 并且支持传入参数 reverse 和 key。

比如我们有一个 tuple 的数组，我们想按照 tuple 的第一个元素进行排序：

l1 = [(1, 2), (0, 1), (3, 10)]

l2 = sorted(l1, key=lambda x: x[0])

# l2 = [(0, 1), (1, 2), (3, 10)]

这里的 key 允许传入一个自定义函数，也可以用自带函数进行比较，比如在一个 string 数组里只想比较小写，可以传入 key=str.lower：

l1 = ["banana", "APPLE", "Watermelon"]
l2 = sorted(l1, key=str.lower)
print(l2)

# ['APPLE', 'banana', 'Watermelon']

多关键字排序时，key 可以返回一个 tuple，Python 会按顺序逐个比较：

points = [(1, 3), (1, 2), (2, 1)]
points.sort(key=lambda p: (p[0], p[1]))
# [(1, 2), (1, 3), (2, 1)]

lambda

你注意到我们在上面使用了 lambda 来定义一个匿名函数，十分方便。如果你熟悉其它语言类似 JS 的话，可以把它理解成一个 callback 函数，参数名一一对应就行。

cmp_to_key

在 Python 3 中，sorted 函数取消了自带的 cmp 函数，如需按自定义比较规则排序，需要借助 functools 库中的 cmp_to_key。不过在绝大多数场景下，优先使用 key 参数（性能更好），cmp_to_key 只在比较逻辑无法用 key 表达时才需要。

比如按照数组元素的绝对值排序，用 key 即可：

arr = [-3, 1, -2, 4]
newarr = sorted(arr, key=abs)
# [1, -2, -3, 4]

只有在复杂的自定义比较时才用 cmp_to_key：

from functools import cmp_to_key

def sortfunc(a, b):
    if abs(a) < abs(b):
        return -1
    elif abs(a) > abs(b):
        return 1
    else:
        return a - b

newarr = sorted(arr, key=cmp_to_key(sortfunc))

set

set 的查找操作复杂度为O(1)，有时候可以替代dict 来存储中间过程。

• add : set 的添加是 add 不是 append
• remove vs discard: 都是删除操作，区别在于 remove 删除不存在的元素会报错，discard 不会。
• union, intersection: 快速获得并集和交集，方便一些去重操作。也可以直接用运算符 | 和 &。
• difference、symmetric_difference：差集和对称差集，运算符为 - 和 ^。

a = {1, 2, 3}
b = {2, 3, 4}
a | b   # {1, 2, 3, 4} 并集
a & b   # {2, 3} 交集
a - b   # {1} 差集
a ^ b   # {1, 4} 对称差集

dict

字典，相当于其它语言中的map, hashtable, hashmap 之类的，读取操作也是O(1) 复杂度。

keys(), values(), items()

这三个方法可以分别获得key, value, {key: value} 的视图。需要注意在 Python 3 中它们返回的是视图对象（view），若需要 list 需要手动 list(...) 转换。

get

get(key, default) 在 key 不存在时返回默认值，而不是抛异常，常用于计数或统计：

d = {}
for c in "hello":
    d[c] = d.get(c, 0) + 1
# {'h': 1, 'e': 1, 'l': 2, 'o': 1}

setdefault

这个函数经常在初始化字典时候使用，如果某个key在字典中存在，返回它的value, 否则插入并返回你给的 default 值。比如在建一个 trie 树的时候：

node = self.root
for char in word:
    node = node.setdefault(char, {})

OrderedDict

OrderedDict 能记录你 key 和 value 插入的顺序，底层其实是一个双向链表加哈希表的实现。我们甚至可以使用move_to_end这样的函数：

>>> d = OrderedDict.fromkeys('abcde')
>>> d.move_to_end('b')
>>> ''.join(d.keys())
'acdeb'
# 放开头
>>> d.move_to_end('b', last=False)
>>> ''.join(d.keys())
'bacde'

注：从 Python 3.7 起，普通的 dict 也已保证插入有序，所以很多时候普通 dict 就能替代 OrderedDict；只有当需要 move_to_end 这类操作时才需要 OrderedDict。

defaultdict

defaultdict 可以很好地来解决一些初始化的问题，比如 value 是一个 list，每次需要判断 key 是否存在的情况。这时我们可以直接定义：

from collections import defaultdict

d = defaultdict(list)

s = [('yellow', 1), ('blue', 2), ('yellow', 3), ('blue', 4), ('red', 1)]
for k, v in s:
    d[k].append(v)
sorted(d.items())
# [('blue', [2, 4]), ('red', [1]), ('yellow', [1, 3])]

常用的 defaultdict(int) 可以用于计数，defaultdict(set) 用于分组去重。

heapq

heapq 就是 python 的 priority queue，heapq[0] 即为堆顶元素。

heapq 的实现是小顶堆，如果需要一个大顶堆，常规的一个做法是把值取负存入，取出时再反转。以下是借助 heapq 来实现 heapsort 的例子：

import heapq

def heapsort(iterable):
    h = []
    for value in iterable:
        heapq.heappush(h, value)
    return [heapq.heappop(h) for i in range(len(h))]

>>> heapsort([1, 3, 5, 7, 9, 2, 4, 6, 8, 0])
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

常用函数：

• heapq.heapify(lst)：将一个 list 原地转换成堆，O(n)
• heapq.heappush(heap, item)：入堆
• heapq.heappop(heap)：弹出堆顶
• heapq.heapreplace(heap, item)：弹出堆顶并压入新元素
• heapq.nlargest(n, iterable) / heapq.nsmallest(n, iterable)：直接取最大/最小的 n 个元素
• 堆中的元素如果是 tuple，会按 tuple 的第一个元素比较，这可用于实现带权重的优先级队列。

bisect

python 自带二分查找的库，在一些不要求实现 binary search，但是借助它能加速的场景下可以直接使用。

import bisect

idx = bisect.bisect_left(a, x, lo=0, hi=len(a))

相似函数还有：

• bisect.bisect_left(a, x)：返回可以插入 x 的最左 index（即第一个 >= x 的位置）
• bisect.bisect_right(a, x) / bisect.bisect：返回可以插入 x 的最右 index（即第一个 > x 的位置）
• bisect.insort_left(a, x) / bisect.insort(a, x)：在有序列表里插入元素并保持有序

Counter

Counter 是 collections 库里的计数器，接受的参数可以是一个 string、一个 list 或者 mapping：

from collections import Counter

>>> c = Counter()                           # a new, empty counter
>>> c = Counter('gallahad')                 # a new counter from an iterable
>>> c = Counter({'red': 4, 'blue': 2})      # a new counter from a mapping
>>> c = Counter(cats=4, dogs=8)             # a new counter from keyword args

• most_common(n)：得到出现次数最多的 n 个数：

>>> Counter('abracadabra').most_common(3)
[('a', 5), ('r', 2), ('b', 2)]

• Counter 之间可以直接用运算符做加减、交集、并集：

a = Counter('aab')
b = Counter('abb')
a + b   # Counter({'a': 3, 'b': 3})
a - b   # Counter({'a': 1})
a & b   # Counter({'a': 1, 'b': 1})  交集（取较小者）
a | b   # Counter({'a': 2, 'b': 2})  并集（取较大者）

strings

ord, chr

ord 返回单个字符的 unicode 码点：

>>> ord('a')
97

chr 则是反向操作：

>>> chr(100)
'd'

在需要构造字母表时，可以用 string 库：

import string
string.ascii_lowercase   # 'abcdefghijklmnopqrstuvwxyz'
string.ascii_uppercase   # 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
string.digits            # '0123456789'

strip

移除 string 前后的字符串，默认移除空格，但也可以传一个字符串，会移除位于首尾、且属于这个字符串的字符：

>>> '   spacious   '.strip()
'spacious'
>>> 'www.example.com'.strip('cmowz.')
'example'

类似的还有 lstrip（去掉左边）和 rstrip（去掉右边）。

split / join

按照某个字符串来切分，返回一个 list，可以传入参数 maxsplit 来限定分离次数。

>>> '1,2,3'.split(',')
['1', '2', '3']
>>> '1,2,3'.split(',', maxsplit=1)
['1', '2,3']
>>> '1,2,,3,'.split(',')
['1', '2', '', '3', '']

join 则是 split 的逆操作，把列表用分隔符连接成字符串：

>>> ','.join(['1', '2', '3'])
'1,2,3'

f-string 格式化

Python 3.6+ 的 f-string 是最推荐的字符串格式化方式：

name = "world"
print(f"hello {name}")
# hello world

x = 3.14159
print(f"{x:.2f}")   # 保留两位小数
# 3.14

常用判断

• s.startswith(prefix) / s.endswith(suffix)：判断开头/结尾
• s.find(sub)：返回子串首次出现的 index，不存在返回 -1
• s.count(sub)：统计子串出现次数
• s.isalpha() / s.isdigit() / s.isalnum()：判断是否全为字母/数字/字母或数字

int/ float

最大, 最小 number

有时候初始化我们需要设定 Math.max() 和 Math.min(), 在 Python 中分别以 float('inf') 和 float('-inf') 表示。Python 3 的 int 支持任意精度，不存在整数溢出的问题。

inf = float('inf')
res = -inf   # 求最大值时初始化
res = inf    # 求最小值时初始化

除法

在 Python 3 中， / 会保留浮点，相当于 float 相除；如果需要整数相除（向下取整），用 //；取余用 %：

>>> 3 / 2
1.5
>>> 3 // 2
1
>>> 3 % 2
1

同时得到商和余数可以用 divmod：

>>> divmod(7, 3)
(2, 1)

次方

在 Python 中为 **：

>>> 2 ** 10
1024

math

math 库提供了很多常用数学函数：

import math

math.gcd(12, 18)          # 6    最大公约数（Python 3.5+）
math.lcm(4, 6)            # 12   最小公倍数（Python 3.9+）
math.comb(5, 2)           # 10   组合数 C(5,2)（Python 3.8+）
math.perm(5, 2)           # 20   排列数 P(5,2)（Python 3.8+）
math.factorial(5)         # 120  阶乘
math.isqrt(8)             # 2    整数平方根（向下取整，Python 3.8+）
math.prod([1, 2, 3, 4])   # 24   列表乘积（Python 3.8+）
math.floor(3.7)           # 3
math.ceil(3.2)            # 4

• 进制转换：bin(10) → '0b1010'，hex(255) → '0xff'，反向用 int('1010', 2) → 10。
• 位运算：a & b（与）、a | b（或）、a ^ b（异或）、~a（取反）、a << k（左移）、a >> k（右移）。

conditions

在 Python 的三项表达式(ternary operation) 与其它语言不太一样：

res = a if condition else b

它表示如果 condition 满足，那么 res = a, 不然 res = b，在类 C 的语言里即为：

res = condition ? a : b;

match 语句（Python 3.10+）

类似其它语言的 switch，用于多分支匹配：

def f(x):
    match x:
        case 0:
            return "zero"
        case 1:
            return "one"
        case _:
            return "other"

海象运算符 := （Python 3.8+）

可以在表达式内部赋值，常用于 while 循环里避免重复计算：

while (line := f.readline()):
    process(line)

# 等价于
line = f.readline()
while line:
    process(line)
    line = f.readline()

any, all

any(), all() 很好理解，就是字面意思，即参数中任何一个为 true 或者全部为 true 则返回 true。经常可以秀一些骚操作，比如 36. Valid Sudoku 这题：

class Solution:
    def isValidSudoku(self, board: List[List[str]]) -> bool:
        row = [[x for x in y if x != '.'] for y in board]
        col = [[x for x in y if x != '.'] for y in zip(*board)]
        pal = [[board[i+m][j+n] for m in range(3) for n in range(3) if board[i+m][j+n] != '.'] for i in (0, 3, 6) for j in (0, 3, 6)]
        return all(len(set(x)) == len(x) for x in (*row, *col, *pal))

itertools

这是 Python 自带的迭代器库，有很多实用的、与遍历、迭代相关的函数。

permutations 排列

list(permutations('ABCD', 2))
# AB AC AD BA BC BD CA CB CD DA DB DC

combinations 组合

list(combinations('ABCD', 2))
# AB AC AD BC BD CD

combinations_with_replacement 允许元素重复使用：

list(combinations_with_replacement('ABC', 2))
# AA AB AC BB BC CC

product 笛卡尔积

list(product('AB', '12'))
# A1 A2 B1 B2

accumulate 累加

list(accumulate([1, 2, 3, 4]))
# [1, 3, 6, 10]  前缀和

传一个函数可以做前缀乘积等：

from operator import mul
list(accumulate([1, 2, 3, 4], mul))
# [1, 2, 6, 24]

groupby 合并

[k for k, g in groupby('AAAABBBCCDAABBB')]
# ['A', 'B', 'C', 'D', 'A', 'B']
[list(g) for k, g in groupby('AAAABBBCCD')]
# [['A','A','A','A'], ['B','B','B'], ['C','C'], ['D']]

注：groupby 只合并相邻的相同元素，通常需要先对数据排序。

其他常用

• chain(*iterables)：把多个可迭代对象串联起来
• zip_longest(a, b, fillvalue=0)：按最长的长度 zip，不足处用 fillvalue 填充
• count(start, step)：无限递增的计数器
• islice(iterable, start, stop)：对迭代器做切片

functools

这个库里有很多高阶函数，包括前面介绍到的 cmp_to_key 以及 reduce，但是比较逆天的有 lru_cache，即 least recently used cache。这个 LRU Cache 是一个常见的面试题，通常用 hashmap 和双向链表来实现，Python 居然直接内置了。

用法即直接作为 decorator 装饰在要 cache 的函数上，以变量值为 key 存储，当反复调用时直接返回计算过的值，例子如下：

lru_cache / cache

from functools import lru_cache

@lru_cache(maxsize=None)
def fib(n):
    if n < 2:
        return n
    return fib(n - 1) + fib(n - 2)

注意：lru_cache 是给带参数的函数用的，当参数是 list、dict 等不可哈希对象时无法直接缓存，需要先转成 tuple 等可哈希类型。此外，写记忆化递归时也可以用 @cache（Python 3.9+，等价于 lru_cache(maxsize=None)）。

lru_cache 实战

https://leetcode.com/problems/stone-game-ii/discuss/345230/Python-DP-Solution

def stoneGameII(self, piles: List[int]) -> int:
    N = len(piles)
    for i in range(N - 2, -1, -1):
        piles[i] += piles[i + 1]
    from functools import lru_cache
    @lru_cache(maxsize=None)
    def dp(i, m):
        if i + 2 * m >= N:
            return piles[i]
        return piles[i] - min(dp(i + x, max(m, x)) for x in range(1, 2 * m + 1))
    return dp(0, 1)

partial

partial 可以固定一个函数的部分参数，生成一个新的函数：

from functools import partial

def add(a, b):
    return a + b

add_one = partial(add, 1)
add_one(2)   # 3

resource

• 这是一个大神用各种 Python trick 解题的 repo，可供娱乐：

  https://github.com/cy69855522/Shortest-LeetCode-Python-Solutions

• 另一个关于各种 python tricks 的总结：

  https://github.com/brennerm/PyTricks

当然 Leetcode 讨论区还会经常见到 StefanPochmann 或者 lee215 这样的大神 Po 一些很秀技的 Python 代码，都是学习范本。