Python 生成式（列表，集合，字典）

列表生成式（list comprehensions）

列表生成式是 Python 语言最重要的特性之一。它使得我们能够用非常简洁的表达式构造一个新的列表，基本形式如下：

[expr for val in collection if condition]

这等效于以下的 for 循环：

result = []
for val in collection:
	if condition:
		result.append(expr)

例如，我们想对下面的字符串列表做些处理，过滤出长度大于 2 的字符串且将它们变为大写：

strings = ['a', 'as', 'bat', 'car', 'dove', 'python']

[x.upper() for x in strings if len(x) > 2]
"""
Out: ['BAT', 'CAR', 'DOVE', 'PYTHON']
"""

我们也可以省略条件。例如，如果我们想生成一个 m × n 的全零矩阵：

[[0] * n for _ in range(m)]

结合列表运算，一个稍微复杂点的例子，生成一个 m × n 的矩阵，矩阵的第一行和第一列元素都为 1，其余元素为 0：

[[1] * n] + [[1] + [0] * (n - 1) for _ in range(m - 1)]

[[1] * n] 生成第一行的全 1 元素，再加上剩下的 m - 1 行，每行第一个元素都为 1。

嵌套列表生成式

我们有如下的姓名列表：

all_names = [['John', 'Emily', 'Michael', 'Mary', 'Steven'],
         ['Maria', 'Juan', 'Javier', 'Natalia', 'Pilar']]

我们想要找出名字中含有多于 2 个 ‘e’ 的，如果借助普通的列表生成式，可以这样写：

names_of_interest = []
for names in all_names:
    enough_es = [name for name in names if name.count('e') >= 2]
    names_of_interest.append(enough_es)

但我们还可以使用嵌套列表生成式：

result = [name for names in all_names for name in names
          if name.count('e') >= 2]
result

这和上面的写法以及下面的双层循环都是等效的：

result = []
for names in all_names:
    for name in names:
        if name.count('e') >= 2:
            result.append(name)

result

也就是说，嵌套列表生成式中的 for 循环顺序和我们在写普通的嵌套 for 循环时是一样的。

虽然我们可以嵌套任意多层循环，但多层的嵌套不利于代码的可读性，通常我们最多使用到两层循环。

再来看两个没有使用条件的嵌套列表生成式：

有如下列表：

tuples = [(1, 2, 3), (4, 5, 6), (7, 8, 9)]

我们想把每个元组中的整数都“展平”成一个整数列表：

flattened = [x for tup in tuples for x in tup]
flattened
"""
Out: [1, 2, 3, 4, 5, 6, 7, 8, 9]
"""

严格来说，下面这个例子才称得上真正的嵌套列表生成式，一个列表生成式里还有另外一个列表生成式：

[[x for x in tup] for tup in tuples]
"""
Out: [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
"""

这和我们上面的例子，生成 m × n 全零矩阵，是类似的。

集合生成式

集合生成式与列表生成式非常类似，只是将中括号换成了花括号：

{expr for val in collection if condition}

例如，对于开头的 strings 列表：

strings = ['a', 'as', 'bat', 'car', 'dove', 'python']

如果我们想要一个仅包含列表中字符串不同长度的集合，那么用集合生成式就是这样的：

unique_lengths = {len(x) for x in strings}
unique_lengths
"""
Out； {1, 2, 3, 4, 6}
"""

我们还可以用一个简单的函数 map 来实现这个任务: set(map(len, strings))

字典生成式

字典生成式的语法如下：

{key-expr : value-expr for val in collection if condition}

将 strings 列表中的字符串和对应位置的索引组合成键值对：

mapping = {val : index for index, val in enumerate(strings)}
mapping
"""
Out: {'a': 0, 'as': 1, 'bat': 2, 'car': 3, 'dove': 4, 'python': 5}
"""

下面这个例子比较有意思：

{i : j for i in range(3) for j in range(3) if i != j}

运行之前你能想到结果是什么吗？键可以当作是确定的，变的则是对应的值。当 j 不断循环时，只要满足 i != j 的条件，值就会被更新。因此最终结果为：

"""
Out: {0: 2, 1: 2, 2: 1}
"""

我们最常用的还是列表生成式，集合生成式、字典生成式与列表生成式类似，作相应推广即可。

References

Python for Data Analysis, 2${}^{\mathrm{n}\mathrm{d}}$ edition. Wes McKinney.