用数学思维理解Comprehension

这一节,我们来分享一个比较有Python特色的用法,叫做Comprehension。简单来说,就是用一个特殊格式的单行for循环,创建各种集合类型的对象。如果这是你第一次接触这个概念,可能会觉得不那么容易掌握。

List comprehension

不过没关系,我们从一个最简单的例子开始。为了创建一个包含1-5的list,我们可以这样:

numbers = [i for i in range(1, 6)]
print(numbers)
# [1, 2, 3, 4, 5]

其中,[i for i in range(1, 6)]就叫做list comprehension。可能你觉得这写起来还不如[1, 2, 3, 4, 5]简单,别着急,我们先来理解它的含义:

一个list comprehension由三部分构成:

  • 一对[]
  • []内部用一个表达式开始,在我们的例子中,也就是i
  • 在表达式后面,跟上一个for循环,然后,再跟上若干可选的forif语句。在我们的例子里,当然只有一个for循环,而没有定义可选的部分;

理解了这个格式之后,该如何理解这个comprehension呢?其实,如果套用数学中关于集合的描述,你一下子就理解了。来看这个问题:如何表达由自然数1-5构成的集合呢?

用数学的方式定义,这个集合是这样的:{ i | i ∈ N, i >= 1, i <= 5 }。怎么样,是不是和list comprehension有点像?只不过:

  • 我们用[]替代了{}
  • 两种定义中的表达式i都表示集合中的元素;
  • i后面,我们用for i in range(1, 6)表达了i ∈ N, i >= 1, i <= 5这个概念;

现在,估计list comprehension看起来就顺眼多了。我们来看一些更复杂的例子。例如,进一步在numbers中只选择偶数,就可以这样:

numbers = [i for i in range(1, 6) if i % 2 == 0]
# [2, 4]

只要我们把if i % 2 == 0这部分也理解成是定义i的作用域就好了。

再来看一个更复杂的例子,我们要把下面这个多维字符串数组,“扁平化”成一个普通的整数数组:

strings = [['1', '2', '3'], ['4', '5', '6'], ['7', '8', '9']]
flatten = [int(i) for items in strings for i in items]
# [1, 2, 3, 4, 5, 6, 7, 8, 9]

在这个例子里,要注意的是两个连续的for循环中变量名的写法,在for items in strings中,我们要做的,是遍历出strings中的每一个子数组。而我们要转换的,应该是这个子数组中的每一个元素,而这就是我们在第二个for循环中要做的事情。因此,list comprehension中,我们要进行变换的,是第二个for循环中的循环变量i

最后,再来看一个组合多个list对象的例子:

xSamples = [2, 5, 8]
ySamples = [2, 6, 8]

points = [(x, y) for x in xSamples for y in ySamples if x != y]
print(points)

在上面的代码里,我们从XY轴的值中,找出了所有不在y=x这条直线上的点。大家可以试着自己分析一下这个list comprehension,应该没有任何问题。现在,你应该就能感受到list comprehension的方便之处了,因为对于最后这个例子,如果我们用普通的代码写出来,会是这样的:

points = []
for x in xSamples:
    for y in ySamples:
        if x != y:
            points.append((x, y))

Set comprehension

除了使用list表示一个集合之外,Python中还有一类更像数学中集合定义方式的类型,叫做set,和list唯一的区别,就是set仅用于表示数据集合,是一个无序类型,这和Swift中的Set类型是类似的。

例如,定义一个包含自然数1-5的set,就可以这样:

numberSet = {x for x in range(1, 6)}
print(numberSet)
# {1, 2, 3, 4, 5}

理解了comprehension的含义之后,上面这段代码就应该没有任何难度了。

Dictionary comprehension

在这一节最后,我们来看dictionary comprehension。这是一个在Python 3中加入的特性,但后来被反向移植到了Python 2.7。实际上,除了comprehension中表达式的部分不同之外,dictionary comprehension没有任何特别之处:

numberDict = {i: str(i) for i in range(1, 6)}
print(numberDict)
# {1: '1', 2: '2', 3: '3', 4: '4', 5: '5'}

另外,通过dictionary comprehension,我们还可以方便的交换它的key-value:

user = {'email': '11@boxue.io', 'workId': 11}
switch = {value: key for key, value in user.items()}
print(switch)
# {'11@boxue.io': 'email', 11: 'workId'}

在上面的例子里,user.items返回的是一个view objects对象,现在没必要太多关注这个细节。只要知道我们可以在for循环里,使用key, value in这样的形式读取到dictionary对应的值就好了。这样,在comprehension表达的部分,我们就能用value:key这样的形式,进行交换了。

推荐阅读更多精彩内容