Python

[Python] 什麼是 List Comprehension

(本文同步發表於計算思維學院)

什麼是 List Comprehension?

List comprehension 比較正式的解釋是「基於一個 iterable 物件中的元素,創造出一個 list」。那什麼又是 iterable 物件?簡單來說就是「可以被 for 迴圈逐一巡覽的物件」。如果以上解釋你還是有聽沒有懂,更直白的說法是:List comprehension 就是用一行程式碼生成一個 list 的寫法。比如你想要產生一個數字序列中的偶數:

list1 = [1, 2, 3, 4, 5, 6, 7, 8, 9]
even_list = [n for n in list1 if n % 2 == 0]
print(even_list)  # [2, 4, 6, 8]

又比如你有一個列表 people,包含所有人的名字跟年齡,你想要找出所有成年人的名字:

people = [("John", 20), ("Jane", 22), ("Judy", 12), ("Jake", 30), ("Jessica", 16)]
adults = [name for name, age in people if age >= 18]
print(adults)  # ['John', 'Jane', 'Jake']

為什麼要使用 List Comprehension?

你可能會問:上面的例子也可以用傳統的 for loop 與 if 判斷式產生,為什麼一定要用 list comprehension? 的確,上面的例子也可以寫成以下這樣:

list1 = [1, 2, 3, 4, 5, 6, 7, 8, 9]
even_list = []
for n in list1:
    if n % 2 == 0:
        even_list.append(n)
print(even_list)  # [2, 4, 6, 8]
people = [("John", 20), ("Jane", 22), ("Judy", 12), ("Jake", 30), ("Jessica", 16)]
adults = []
for name, age in people:
    if age >= 18:
        adults.append(name)
print(adults)  # ['John', 'Jane', 'Jake']

但是你可以看出來,傳統的寫法程式碼會比較長,雖然程式碼的主要判斷標準還是以可讀性為主,比較長不一定比較差,但 list comprehension 已經是 Python 社群中廣為接受的撰寫風格(換句話說就是比較 Pythonic),所以在不影響可讀性的狀況下,list comprehension 通常可以較簡潔地表達程式邏輯。

另外,雖然兩種寫法在理論上的時間/空間複雜度是相同的,但 list comprehension 在實務上的速度可能會稍快,因為其底層實作有經過一些最佳化。比如以下是產生一百萬個數字的平方列表的兩種方式:

import timeit
# 寫法 1: list comprehension
list_comprehension_code = """
squared = [x**2 for x in range(1, 1000001)]
"""
# 寫法 2: for loop
for_loop_code = """
squared = []
for x in range(1, 1000001):
    squared.append(x**2)
"""
# 使用 timeit 測量代碼執行時間
lc_time = timeit.timeit(list_comprehension_code, number=100)  # 執行 100 次
fl_time = timeit.timeit(for_loop_code, number=100)  # 執行 100 次
print(f"List comprehension time: {lc_time}")  # 21.996
print(f"For loop time: {fl_time}")  # 25.540

這兩段程式碼在我的筆電上執行 100 次之後,for loop 的寫法會慢 16% 左右。

使用 List Comprehension 的時機

要怎麼選擇何時使用 list comprehension 或 for loop? 其實正確的答案是:你開心(或團隊開心)就好。一般的建議是:

  • 當程式碼邏輯的焦點是產生新的 list, 可以考慮使用 list comprehension
  • 當程式碼邏輯的焦點是處理原本的序列,不需要產生新的 list, 那 for loop 足矣
  • 如果產生新的 list 的判斷邏輯相當複雜、冗長,使用 list comprehension 會降低可讀性,那還是使用 for loop 就好

但是這只是一般性的指導原則,如果你是在團隊中,那其實就跟團隊的 coding style 一致,你的同事在 code review 時沒有意見即可。

所以也有 Dict/Set Comprehension 囉?

是的,如果你基於一個 iterable 物件中的元素,創造出一個 dict (set),或是用一行程式碼生成一個 dict (set),就叫做 dict (set) comprehension。下面是一個 dict comprehension 的例子:

student_scores = {
    "John": 85,
    "Jane": 90,
    "Judy": 55,
    "Jake": 78,
    "Jessica": 59
}
# 找出考試 80 分以上的人名及分數
top_students = {student: score for student, score in student_scores.items() if score >= 80}
print(top_students )  # {'John': 85, 'Jane': 90}

而以下是一個 set comprehension 的例子:

course1_students = ["John", "Jane", "Judy", "Jake", "Jessica"]
course2_students = ["Jane", "Jake", "Sam", "Sue", "Sandra"]
# 找出兩堂課都有上的學生名字
common_students = {name for name in course1_students if name in course2_students}
print(common_students)  # {'Jake', 'Jane'}

想系統化學習更多 Python 資料型態與進階實務觀念,可以參考 Python 練功坊: 50 道精選練習題助你掌握 Python 實務觀念