Python 컴프리헨션 및 표현식 정리

Credit: https://www.linkedin.com/pulse/python-tips-list-comprehensions-hugo-tota

Python으로 알고리즘을 공부할 때 놓칠 수 없는 개념들 중 하나가 바로 컴프리헨션(comprehension)과 표현식(expression)입니다. Python에서는 반복문과 조건문을 간결하게 표현할 수 있는 list comprehension, generator expression, conditional expression(삼항 연산자) 등을 제공하여, 적절하게 사용 시 성능(실행속도)과 간결함을 둘 다 잡을 수 있습니다.

이번 포스트에서는 각종 컴프리헨션 및 표현식의 개념과 사용법을 정리합니다.

아울러, 이들에 람다(lambda)식까지 적용한 예제도 한번 살펴보시겠습니다.

---

목차:

1. List Comprehensions
2. Set Comprehensions
3. Dictionary Comprehensions
4. Generator Expressions
5. Conditional Expressions
6. Nested Loops in Comprehensions/Expressions
7. Lambda 적용하기
8. 더 나아가기 - 안티패턴

---

1. List Comprehensions

리스트 컴프리헨션(List Comprehension)은 기존 리스트나 이터러블 객체(반복가능 객체)에서 새로운 리스트를 간결하고 효율적으로 생성할 수 있게 해주는 문법입니다. for문 뒤에 if를 붙임으로써 필터링 기능도 수행할 수 있습니다.

[표현식 for 변수 in 이터러블 if 조건]

특징:

• 필터링과 변환을 동시에 처리
• 기존 for문보다 빠르고, 코드도 짧음 ==> 성능 & 가독성 개선
• 조건부 표현식 (conditional expressions)와 함께 사용 가능 ==> 필터링 기능

예제 1-0: range()로 0부터 9까지의 리스트 생성

# 기존 for문
numbers_list = []
for x in range(10):
    numbers_list.append(x)

# List comprehension 적용
numbers_list = [x for x in range(10)]

# 출력
print(numbers_list) # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

예제 1-1: 문자열 형태의 입력을 정수형으로 변환 => 리스트로 생성

numbers_str = "1 2 43 2 325 25 5 354 36 47"
numbers_list = [x for x in map(int, numbers_str.split())]

# 출력
print(numbers_list)  # [1, 2, 43, 2, 325, 25, 5, 354, 36, 47]

예제 1-2: 단어 길이 구하기

sentence = "Python is fun"
word_lengths = [len(word) for word in sentence.split()]
print(word_lengths)  # [6, 2, 3]

예제 1-3: 1부터 10까지 중 짝수만 필터링 후 각각의 제곱수 출력

squares = [x**2 for x in range(1, 11) if x % 2 == 0]
print(squares)  # [4, 16, 36, 64, 100]

예제 1-4: (요소, 제곱)의 튜플의 리스트 생성

numbers = [1, 2, 3, 4]
tuple_list = [(x, x**2) for x in numbers]
print(tuple_list)  # [(1, 1), (2, 4), (3, 9), (4, 16)]

예제 1-5: 2차원 리스트 생성 (중첩된 리스트 컴프리헨션)

# 3개의 행(row)와 2개의 열(column)을 가지는 2차원 리스트
# 즉 2x3의 2차원 리스트 - 헷갈리신다면, 천천히 보세요, 6. Nest Loops 항목에서 더 헷갈리실 수 있습니다

matrix_a = [[None]*2 for i in range(3)]
matrix_b = [[None for j in range(2)] for i in range(3)]

print(matrix_a) # [[None, None], [None, None], [None, None]]
print(matrix_b) # [[None, None], [None, None], [None, None]]

예제 1-6: 구구단 2차원 리스트 생성 (중첩된 리스트 컴프리헨션)

multiplication_table = [[i * j for j in range(1, 10)] for i in range(2, 10)]
print(row)

 

2. Set Comprehensions

집합(set)을 생성하는 컴프리헨션. 리스트 컴프리헨션과 같지만 중괄호 {}로 감싸고, 중복 없이 요소를 수집합니다.

{표현식 for 변수 in 이터러블 if 조건}

특징:

• 중복 제거된 결과가 자동 생성됨
• 데이터 필터링 후 중복 없는 결과 필요할 때 유용

예제 2-1: 문자열에서 모음 추출하여 집합 생성

text = "Hello World"
vowels = {char.lower() for char in text if char.lower() in 'aeiou'}
print(vowels)  # {'e', 'o'}

예제 2-2: 1부터 21까지, 5의 배수의 집합 생성

multiples_of_5 = {x for x in range(1, 21) if x % 5 == 0}
print(multiples_of_5)  # {5, 10, 15, 20}

 

3. Generator Expressions

제너레이터 표현식은 리스트 컴프리헨션과 유사하지만 ()로 감쌉니다. 값을 하나씩 생성하므로 (lazy evaluation), 이를 잘 활용하면 메모리를 덜 쓸 수 있습니다.

(표현식 for 변수 in 이터러블 if 조건)

특징:

• 메모리에 모든 요소를 올리지 않고, 하나씩 lazy evaluation
• sum(), max(), any(), all()과 함께 자주 사용합니다.
• 값을 단 한 번만 꺼낼 수 있습니다. 즉 재사용 불가

예제 3-1: 1부터 11까지의 제곱을 프린트

squares_gen = (x**2 for x in range(1, 11))
for square in squares_gen:
    print(square)

예제 3-2: 3의 배수의 합 (제너레이터 표현식)

sum_of_multiples = sum(x for x in range(1, 101) if x % 3 == 0)
print(sum_of_multiples)  # 1683

예제 3-3: 2차원 리스트에서 조건 만족 요소 개수 세기

matrix = [
    [3, 5, 7],
    [2, 6, 4],
    [8, 1, 9]
]
count_ge_5 = sum(1 for row in matrix for item in row if item >= 5)
print(count_ge_5)  # 5

예제 3-4: 리스트 합과 평균 계산하기

numbers = [10, 20, 30, 40, 50]
total = sum(x for x in numbers)
average = total / len(numbers)
print(total, average)  # 150 30.0

 

4. Conditional Expressions

컴프리헨션 내부에서 사용할 수 있는 삼항 연산자입니다. 나오는 값 자체를 조건에 따라 바꿀 때 사용됩니다. 예제 1-3.에서 보여드린 필터링과는 구별되는 개념입니다. for문을 사용하지 않는 표현식도 존재합니다 (예제 4-1).

[참일때의값 if 조건 else 거짓일때의값 for 변수 in 이터러블]

특징:

• 나오는 값 자체를 조건에 따라 변환함 ==> 필터링과는 다름!
• 필터링과 혼동하지 않도록 주의

예제 4-1: 경우에 따른 문자열 출력

def foo():
    return True

foo_result = 'YES' if foo() else 'NO'

print(foo_result) # YES

예제 4-2: 조건부 표현식을 사용한 리스트 변환

numbers = [1, 2, 3, 4, 5, 6]
even_or_zero = [x if x % 2 == 0 else 0 for x in numbers]
print(even_or_zero)  # [0, 2, 0, 4, 0, 6]

예제 4-3: 짝수/홀수 매핑 리스트 생성하기

numbers = [1, 2, 3, 4, 5]
parity = ["Even" if x % 2 == 0 else "Odd" for x in numbers]
print(parity)  # ['Odd', 'Even', 'Odd', 'Even', 'Odd']

 

5. Dictionary Comprehensions

딕셔너리 컴프리헨션은 {key: value} 형태로 딕셔너리를 빠르게 생성합니다. Set처럼 { }를 사용하지만, 얘는 key: value의 조합이므로, :가 들어감으로써 set comprehensions와 구별됩니다.

{key_expr: value_expr for 변수 in 이터러블 if 조건}

특징:

• 키-값 쌍을 동적으로 생성할 때 유용
• 특정 규칙으로 매핑된 딕셔너리를 만들기 좋음

예제 5-1: 숫자와 제곱의 딕셔너리 만들기

numbers = [1, 2, 3, 4, 5]
square_dict = {x: x**2 for x in numbers}
print(square_dict)  # {1: 1, 2: 4, 3: 9, 4: 16, 5: 25}

예제 5-2: 문자열의 첫 글자에 따른 알파벳 순서 번호 딕셔너리

words = ['apple', 'banana', 'cherry']
alpha_index = {word: ord(word[0].lower()) - ord('a') + 1 for word in words}
print(alpha_index)  # {'apple': 1, 'banana': 2, 'cherry': 3}

예제 5-3: 그래프에서, 각 노드를 딕셔너리로 저장 (백준 1991 "트리 순회" 참조)

이 예제는 generator expression과 dictionary comprehension을 중첩한 경우입니다.

# 입력 예시 - 백준 1991 "트리 순회" 참조
"""
A B C
B D .
C E F
E . .
F . G
D . .
G . .
"""
tree = {root: (left, right) for root, left, right in (input().split() for _ in range(n))}
print(tree) # {'A': ('B', 'C'), 'B': ('D', '.'), 'C': ('E', 'F'), 'E': ('.', '.'), 'F': ('.', 'G'), 'D': ('.', '.'), 'G': ('.', '.')}

6. Nested Loops in Comprehensions/Expressions

Nested loops, 즉, 2개 이상의 for문이 중첩되는 형태입니다. 다차원 리스트 등을 만들거나 여러 리스트의 조합을 계산할 때 유용합니다. Nested comprehensions/expressions와는 구별되는 개념입니다.

[표현식 for 변수1 in 이터러블1 for 변수2 in 이터러블2 ...]

특징:

• 이중 반복문을 한 줄로 표현 가능
• 다차원 구조(행렬 등)에 적합

예제 6-1: 리스트 조합으로 튜플 생성

pairs = [(x, y) for x in [1, 2, 3] for y in ['a', 'b']]
print(pairs) # [(1, 'a'), (1, 'b'), (2, 'a'), (2, 'b'), (3, 'a'), (3, 'b')]

예제 6-2: 리스트 조합으로 각 합의 리스트 생성

list1 = [1, 2]
list2 = [10, 20, 30]
combination_sums = [x + y for x in list1 for y in list2]
print(combination_sums)  # [11, 21, 31, 12, 22, 32]

예제 6-3: 다차원 리스트를 단차원으로 펴기 (flattening)

my_lists = [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]]
flat = [fooooooo for sublist1 in my_lists for sublist2 in sublist1 for fooooooo in sublist2]
print(flat) # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]

예제 6-4: 다차원 리스트 내에 0이라는 값이 존재하면 "YES", 없으면 "NO" 표시

# 리스트 내에 0이라는 값이 하나라도 존재하면 "YES", 이외에는 "NO"를 프린트
block = [[[0, -1, 0, 0, 0], [-1, -1, 0, 1, 1], [0, 0, 0, 1, 1]]]
print("YES" if any(0 in col for row in block for col in row) else "NO") # YES

# 적당한 indentation 적용 시
print("YES" if any( 0 in col 
                    for row in block 
                    for col in row)  else "NO")

※ 위 1. List Comprehensions의 2차원 리스트 생성과 비교하여, for문을 어떤 순서로 읽어야 할지를 중점적으로 보세요.

 

7. Lambda 적용하기

Lambda란 쉽게 말하면 '익명 함수'로, 단일 문장(1줄)의 코드로 작성됩니다. 함수 내부에는 return문이 포함되지 않지만, 식 자체로부터 나오는 값을 반환합니다.

위 comprehensions 및 expressions에 lambda를 적용하여, 기존의 이터러블 객체로부터 새로운 리스트 등을 간결하게 생성할 수 있습니다.

# 아래는 x+y 값을 반환하는 lambda식입니다.
add = lambda x, y: x + y

예제 7-1: 각 요소의 제곱의 리스트 컴프리헨션 ==> 프린트

numbers = [1, 2, 3, 4, 5]
square = lambda x: x ** 2
squares = [square(x) for x in numbers]
print(squares)  # [1, 4, 9, 16, 25]

예제 7-2: 짝수인 요소만 추려서 해당 요소들의 제곱의 제너레이터 표현식 => 합계 

square = lambda x: x ** 2
sum_of_even_squares = sum(square(x) for x in range(1, 11) if x % 2 == 0)
print(sum_of_even_squares)  # 220 (4 + 16 + 36 + 64 + 100)

 

8. 더 나아가기 - 안티패턴

서두에 성능와 간결함을 잡을 수 있다고 언급 드렸으나, 식이 지나치게 복잡해지는 경우에는 오히려 일반적인 반복문보다 가독성이 나빠지게 됩니다. 예를 들어, 아래와 같이 짜느니 차라리 기존의 방식대로 작성하는 것이 좋을 수 있습니다.

# 경고: 아래 코드는 고의로 난해하게, 유지보수하기 힘들도록 작성한 안티패턴 예시임.
bad_sample = {f"key_{i}": ([[[x if x % 2 == 0 else -x for x in range(j, j + 3)] for j in range(k, k + 2)] for k in range(i, i + 3)] if i % 2 == 0 else {f"subkey_{k}": {f"inner_{n}": n if n % 3 == 0 else n * 2 for n in range(k, k + 3)} for k in range(i, i + 3)}) for i in range(1, 4)}

import pprint
pprint.pprint(bad_sample)

""" 실행 결과:
{'key_1': {'subkey_1': {'inner_1': 2, 'inner_2': 4, 'inner_3': 3},
           'subkey_2': {'inner_2': 4, 'inner_3': 3, 'inner_4': 8},
           'subkey_3': {'inner_3': 3, 'inner_4': 8, 'inner_5': 10}},
 'key_2': [[[2, -3, 4], [-3, 4, -5]],
           [[-3, 4, -5], [4, -5, 6]],
           [[4, -5, 6], [-5, 6, -7]]],
 'key_3': {'subkey_3': {'inner_3': 3, 'inner_4': 8, 'inner_5': 10},
           'subkey_4': {'inner_4': 8, 'inner_5': 10, 'inner_6': 6},
           'subkey_5': {'inner_5': 10, 'inner_6': 6, 'inner_7': 14}}}
"""

---

참고자료:

https://wikidocs.net/138214