When to use list comprehension¶

When to Use a List Comprehension in Python – Real Python

难度: 3

时长: 30 min

Transforming Lists in Python¶

`for` Loops¶

>>> squares = []
>>> for number in range(10):
...     squares.append(number * number)
...
>>> squares
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

With `map` Objects¶

>>> prices = [1.09, 23.56, 57.84, 4.56, 6.78]
>>> TAX_RATE = .08
>>> def get_price_with_tax(price):
...     return price * (1 + TAX_RATE)
...

>>> final_prices = map(get_price_with_tax, prices)
>>> final_prices
<map object at 0x7f34da341f90>

>>> list(final_prices)
[1.1772000000000002, 25.4448, 62.467200000000005, 4.9248, 7.322400000000001]

With List Comprehension¶

new_list = [expression for member in iterable]

>>> squares = [number * number for number in range(10)]
>>> squares
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

expression is the member itself, a call to a method, or any other valid expression that returns a value. In the example above, the expression number * number is the square of the member value.
member is the object or value in the list or iterable. In the example above, the member value is number.
iterable is a list, set, sequence, generator, or any other object that can return its elements one at a time. In the example above, the iterable is range(10).

>>> prices = [1.09, 23.56, 57.84, 4.56, 6.78]
>>> TAX_RATE = .08
>>> def get_price_with_tax(price):
...     return price * (1 + TAX_RATE)
...

>>> final_prices = [get_price_with_tax(price) for price in prices]
>>> final_prices
[1.1772000000000002, 25.4448, 62.467200000000005, 4.9248, 7.322400000000001]

Supercharging Your Python List Comprehensions¶

Filter Values From a List¶

The most common way to add conditional logic to a list comprehension is to add a conditional to the end of the expression.

new_list = [expression for member in iterable if conditional]

>>> sentence = "the rocket came back from mars"
>>> [char for char in sentence if char in "aeiou"]
['e', 'o', 'e', 'a', 'e', 'a', 'o', 'a']

>>> sentence = (
...     "The rocket, who was named Ted, came back "
...     "from Mars because he missed his friends."
... )
>>> def is_consonant(letter):
...     vowels = "aeiou"
...     return letter.isalpha() and letter.lower() not in vowels
...

>>> [char for char in sentence if is_consonant(char)]
['T', 'h', 'r', 'c', 'k', 't', 'w', 'h', 'w', 's', 'n', 'm', 'd',
 'T', 'd', 'c', 'm', 'b', 'c', 'k', 'f', 'r', 'm', 'M', 'r', 's', 'b',
 'c', 's', 'h', 'm', 's', 's', 'd', 'h', 's', 'f', 'r', 'n', 'd', 's']

new_list = [true_expr if conditional else false_expr for member in iterable]

>>> original_prices = [1.25, -9.45, 10.22, 3.78, -5.92, 1.16]
>>> [price if price > 0 else 0 for price in original_prices]
[1.25, 0, 10.22, 3.78, 0, 1.16]

>>> def get_price(price):
...     return price if price > 0 else 0
...

>>> [get_price(price) for price in original_prices]
[1.25, 0, 10.22, 3.78, 0, 1.16]

Remove Duplicates With Set and Dictionary Comprehensions¶

>>> quote = "life, uh, finds a way"
>>> {char for char in quote if char in "aeiou"}
{'a', 'e', 'u', 'i'}

Dictionary comprehension

>>> {number: number * number for number in range(10)}
{0: 0, 1: 1, 2: 4, 3: 9, 4: 16, 5: 25, 6: 36, 7: 49, 8: 64, 9: 81}

Assign Values With the Walrus Operator¶

>>> import random
>>> def get_weather_data():
...     return random.randrange(90, 110)
...

>>> [temp for _ in range(20) if (temp := get_weather_data()) >= 100]
[107, 102, 109, 104, 107, 109, 108, 101, 104]

Deciding When Not to Use a List Comprehension¶

Watch Out for Nested Comprehensions¶

>>> cities = ["Austin", "Tacoma", "Topeka", "Sacramento", "Charlotte"]
>>> {city: [0 for _ in range(7)] for city in cities}
{
    'Austin': [0, 0, 0, 0, 0, 0, 0],
    'Tacoma': [0, 0, 0, 0, 0, 0, 0],
    'Topeka': [0, 0, 0, 0, 0, 0, 0],
    'Sacramento': [0, 0, 0, 0, 0, 0, 0],
    'Charlotte': [0, 0, 0, 0, 0, 0, 0]
}

>>> [[number for number in range(5)] for _ in range(6)]
[
    [0, 1, 2, 3, 4],
    [0, 1, 2, 3, 4],
    [0, 1, 2, 3, 4],
    [0, 1, 2, 3, 4],
    [0, 1, 2, 3, 4],
    [0, 1, 2, 3, 4]
]

2 维 list 的 comprehension:

matrix = [
...     [0, 0, 0],
...     [1, 1, 1],
...     [2, 2, 2],
... ]

>>> [number for row in matrix for number in row]
[0, 0, 0, 1, 1, 1, 2, 2, 2]

不是十分容易理解，此时可能用原始2维 list 的 for 循环会更好：

>>> matrix = [
...     [0, 0, 0],
...     [1, 1, 1],
...     [2, 2, 2],
... ]
>>> flat = []
>>> for row in matrix:
...     for number in row:
...         flat.append(number)
...
>>> flat
[0, 0, 0, 1, 1, 1, 2, 2, 2]

Choose Generators for Large Datasets¶

If you sum the first billion squares with a generator, then your program will likely run for a while, but it shouldn’t cause your computer to freeze:

>>> sum(number * number for number in range(1_000_000_000))
333333332833333333500000000

The example above still requires a lot of work, but it performs the operations lazily. Because of lazy evaluation, your code only calculates values when they’re explicitly requested.

The map() function also operates lazily, meaning memory won’t be an issue if you choose to use it in this case:

>>> sum(map(lambda number: number * number, range(1_000_000_000)))
333333332833333333500000000

Profile to Optimize Performance¶

>>> import random
>>> import timeit
>>> TAX_RATE = .08
>>> PRICES = [random.randrange(100) for _ in range(100_000)]
>>> def get_price(price):
...     return price * (1 + TAX_RATE)
...
>>> def get_prices_with_map():
...     return list(map(get_price, PRICES))
...
>>> def get_prices_with_comprehension():
...     return [get_price(price) for price in PRICES]
...
>>> def get_prices_with_loop():
...     prices = []
...     for price in PRICES:
...         prices.append(get_price(price))
...     return prices
...

>>> timeit.timeit(get_prices_with_map, number=100)
2.0554370979998566

>>> timeit.timeit(get_prices_with_comprehension, number=100)
2.3982384680002724

>>> timeit.timeit(get_prices_with_loop, number=100)
3.0531821520007725