[数据结构与算法]Effective Python 第27、28、29条：列表、集合推导式用法总结和建议

python里可以根据某个序列或可迭代对象派生出一个新的列表。用这种写法的表达式，叫列表推导式。列表推导式可以大大简化代码。

一般推导式

如下，把原来的列表每个元素平方放到一个新列表里：

a = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

b = [x**2 for x in a]
print(b)
# [1, 4, 9, 16, 25, 36, 49, 64, 81, 100]

字典和集合也有相应的推导式，如下：

a = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

b_dict = {x: x**2 for x in a}
c_set = {x**2 for x in a}
print(b_dict)
print(c_set)
# {1: 1, 2: 4, 3: 9, 4: 16, 5: 25, 6: 36, 7: 49, 8: 64, 9: 81, 10: 100}
# {64, 1, 4, 36, 100, 9, 16, 49, 81, 25}

带条件的推导式

如下，把原列表的偶数元素过滤出来，平方后放在新的列表里：

a = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

b_dict = {x: x**2 for x in a if x % 2 == 0}
print(b_dict)
# {2: 4, 4: 16, 6: 36, 8: 64, 10: 100}

另外，还支持if else分开处理，如下：

a = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

b = [x**2 if x % 2 == 0 else x + 1 for x in a]
print(b)
# [2, 4, 4, 16, 6, 36, 8, 64, 10, 100]

还支持多个if条件，如下：

a = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

b = [x for x in a if x > 4 if x % 2 == 0]
print(b)  # [6, 8, 10]

多层推导

列表推导式还支持多层循环，如下，把二维列表展开：

a = [[1, 2, 3], [4, 5, 6], [7, 8, 9, 10]]

b = [x for row in a for x in row]
print(b)
# [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

建议：不要写太多循环，最多两层，否则不易理解，超过两层还是for循环展开。

结合赋值表达式简化代码

python3.8之后引入了赋值表达式，详情可见赋值表达式使用。

推导list、dict、set时，经常要在多个地方用到同一个计算结果。例如，在库存信息里，有各个商品的数量，需要检查商品是否够8个，8个为一批，才可以发货，如下：

stock = {  # 各个商品的数量
    'a': 125,
    'b': 35,
    'c': 8,
    'd': 24,
}

order = ['a', 'c', 'e']  # 订单中的各个商品


def get_batches(count, size):
    return count // size


# 使用字典推导式计算订单每个商品够几批了
found = {name: get_batches(stock.get(name, 0), 8)
         for name in order
         if get_batches(stock.get(name, 0), 8)}
print(found)
# {'a': 15, 'c': 1}

可以发现get_batches使用了两次。这样写是代码比较混乱，可以使用赋值表达式来改进，如下：

found = {name: batches
         for name in order
         if (batches := get_batches(stock.get(name, 0), 8))}
print(found)
# {'a': 15, 'c': 1}

batches := get_batches(stock.get(name, 0), 8)这个就是赋值表达式，能够从stock字典汇总查到对应产品一共有几批，并把这个批数放到batches变量里。这样，推导这个产品所对应的批数就不用在使用get_batches函数获取了。

注意不要把赋值表达式用在推导式里，if条件用他，这样会报错，如下：

result = {name: (tenth := count // 10)
          for name, count in stock.items() if tenth > 0}
# NameError: name 'tenth' is not defined

变量泄露

普通的for循环会存在变量泄露的情况，如下：

a = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

for i in a:
    pass

# 循环里的变量i在外面也能用
print(i)  # 10

但是，推导语句中的for循环所使用的循环变量，是不会像刚才那样泄露到外面的，如下：

a = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

b = [i**2 for i in a]

# 推导式循环里的变量i不会泄露在外面
print(i)  # NameError: name 'i' is not defined