[Python] 리스트 함축(list comprehension)을 시작부터 끝까지 이해하여보자.

파이썬에서의 list를 야무지게 사용하는 방법 중 하나는 리스트 함축(list comprehension)이다. 이 또한 파이썬 이외에 언어에서는 찾아볼 수 없는, 독특한 문법 중 하나이다.

 

list를 새로 생성할 때 매우 유용하게 사용할 수 있으며, 특히 코드의 가독성을 높이고 코드의 길이를 짧게 할 때 유용하게 사용할 수 있는 테크닉이다.

 

이번 포스팅에서는 리스트 함축에 대하여 알아보고 이를 극한까지 활용하여보자.

list comprehension에 대하여 알아보자.

 

---

1. list comprehension이란?

파이썬 공식 문서에는 리스트 함축(list comprehension)에 대한 정의를 다음과 같이 하고 있다.

List comprehension은 list를 생성하는데 하나의 간결한 방법을 제공하는 방법입니다. 다른 sequence 혹은 iterable에서 일정한 연산을 수행한 결과들의 list을 만들거나, 일정 조건을 만족하는 list를 생성하는데 자주 사용됩니다.

 

위 정의대로 일정 조건을 만족하는 새로운 list 혹은 다른 sequence(range, 다른 list, dictionary...)에서 일정한 연산을 취한 후 이에 대한 결과를 담는데 유용하게 사용할 수 있는 방법이다.

 

다음 예제를 보자. 만약 1부터 9까지의 제곱수를 담은 list를 생성하려면 다음과 같이 작성할 수 있다.

 

some_list = []
for i in range(10):
    some_list.append(i*i)
 
print(some_list)
# [1, 4, 9, 16, 25, 36, 49, 64, 81]

 

흐음... 간단한 연산을 하는데도 5줄이나 필요하다. 만약 람다 함수에 대하여 알고 있다면, 다음과 같이 map과 익명(람다) 함수를 이용하여 생성할 수 있다.

 

some_list = list(map(lambda x : x*x, range(10)))
 
print(some_list)
# [1, 4, 9, 16, 25, 36, 49, 64, 81]

 

그러나, 가독성이 떨어진다는 단점이 존재한다. 만약 list comprehension을 활용한다면, 다음과 같이 한 줄로 아주 간단하게 작성할 수 있다!

 

print([x*x for x in range(10)])
# [1, 4, 9, 16, 25, 36, 49, 64, 81]

 

위의 코드의 경우 람다 함수와 map을 활용한 경우보다 가독성이 높아지고 유지보수하기 쉬워진 형태가 되었다.

---

list comprehension의 경우 x를 제곱하는 것 이외에 여러 가지 일을 할 수 있다.

 

예를 들어, 다음과 같이 x에 여러 메서드를 취할 수 있다. list안의 원소들을 대문자로 바꾸고 싶은 경우 다음과 같이 list comprehension을 이용하여 작성할 수 있다.

 

print([x.upper() for x in ["a", "b", "c"]])
# ['A', 'B', 'C']

 

혹은 x를 함수의 인자로 넘길 수 있다. 

 

print([int(x) for x in ["1", "2", "3"]])
# [1, 2, 3]

 

혹은 다음과 같이 x에 대하여 여러 연산들 및 함수들을 복합적으로 적용하여도 된다. 

 

print([[int(x), float(x)] for x in ["1", "2", "3"]])
# [[1, 1.0], [2, 2.0], [3, 3.0]]
 
from math import pi
print([str(round(pi, i)) for i in range(1, 6)])
# ['3.1', '3.14', '3.142', '3.1416', '3.14159']

 

list comprehension을 사용하면서 창의력을 발휘하여 보자.

---

list comprehension의 정의에서 일정 조건을 만족하는 list를 생성할 수도 있다고 하였다. 만약 1부터 9까지의 list에서 홀수만 있는 list를 새로 만들고 싶다면, 다음과 같이 작성할 수 있다.

 

some_list = [1, 2, 3, 4, 5, 6, 7, 8, 9]
new_list = []
for x in some_list:
    if x%2 == 1:
        new_list.append(x)
 
print(new_list)
# [1, 3, 5, 7, 9]

 

흠... 너무 길다. 이를 list comprehension을 이용하여 한 줄로 줄여보자.

 

some_list = [1, 2, 3, 4, 5, 6, 7, 8, 9]
print([x for x in some_list if x%2 == 1])
# [1, 3, 5, 7, 9]

 

매우 간단하게 한 줄로 작성할 수 있다. 다만, list comprehension을 사용할 경우 if의 위치는 맨 뒤에 위치하여야 한다. 만약 앞에 위치할 경우 if에 대응하는 else가 반드시 존재해야 한다. 만약 없을 경우 SyntaxError로 에러 메시지가 출력된다.

 

some_list = [1, 2, 3, 4, 5, 6, 7, 8, 9]
print([x if x%2 == 1 else 0 for x in some_list])
# [1, 0, 3, 0, 5, 0, 7, 0, 9]
print([x for x in some_list if x%2 == 1 else 0])
# 오류!

 

이를 정리하면 다음과 같이 쓸 수 있다. 만약 else가 필요하지 않을 경우 else를 제외할 수 있다. 다른 언어에서의 ternary operator(a ? b : c)과 같이 작성하면 된다.

[expression1 for item in list if condition1 else expression2]

은 다음과 같다.

for item in list:
    if condition:
        expression1
    else:
        expression2

---

2. list comprehension의 활용

list comprehension은 하나의 for문 이외에도 여러 개의 for문을 사용할 수 있다.

 

만약 다음과 같이 두 개의 list에서 같은 원소가 아닐 경우 두 list를 합치는 경우를 생각하여보자.

 

x = [1, 2, 3, 4]
y = [3, 4, 5, 6]
new_list = []
for num1 in x:
    for num2 in y:
        if num1 == num2:
            new_list.append(num1)
print(new_list)
# [3, 4]

 

이 경우 중첩된 for문 두 개를 활용하여 구할 수 있다. 이를 for문 두 개를 활용한 list comprehension을 이용하여 나타내면 다음과 같이 사용할 수 있다. for문을 잘 따라가 보면 이해할 수 있다.

 

x = [1, 2, 3, 4]
y = [3, 4, 5, 6]
new_list = [num1 for num1 in x for num2 in y if num1 == num2]
print(new_list)
# [3, 4]

 

여기에 사용된 list comprehension은 중첩된 for문 두 개를 사용한 코드와 같은 list를 반환한다.

 

다른 예제를 보자. list comprehension을 이용하여 두 개의 list가 중첩되어 있는 list를 하나의 list로 만들고 싶다면 다음과 같이 작성할 수 있다.

 

num_2d = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
print([x for item1 in num_2d for x in item1])
# [1, 2, 3, 4, 5, 6, 7, 8, 9]

 

코드를 따라가보자. 첫 번째 for문(for item1 in num_2d)을 통해 item1에 각각 list [1, 2, 3], [4, 5, 6], [7, 8, 9]가 담기고, 이후 두 번째 for문(for x in item1)을 통해 x에 각 list의 원소들을 담는다.

 

이 경우 for문의 순서에 주의하여야 한다. 만약 두 for문의 위치를 바꿀 경우 item1을 찾지 못하여 NameError 에러 메시지를 볼 수 있다.

 

num_2d = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
print([x for x in item1 for item1 in num_2d])
# Traceback (most recent call last):
#   File "main.py", line 2, in <module>
#     print([x for x in item1 for item1 in num_2d])
# NameError: name 'item1' is not defined.

 

for문이 굳이 2개가 아니어도 된다. 극단적인 경우지만 다음과 같이 for문의 개수에는 제한이 없다.

 

l = [[[[[[1]]]]]]
new_list = [float(f) for a in l for b in a for c in b for d in c for e in d for f in e]
print(new_list)
# [1.0]
"""        
           l a b c d e f
           ↓ ↓ ↓ ↓ ↓ ↓ ↓
new_list = [ [ [ [ [ [ 1 ] ] ] ] ] ]
과 같은 순서로 이루어진다.
"""
 

 

사실 3개 이상의 for문 list comprehension은 알고리즘을 공부할 때 이외에는 다루는 경우가 정말 없기에, '이런 게 있구나'하고 넘어가도 된다.

---

만약 list comprehension에서 if...else 대신 if...elif...else를 사용할 때에는 어쩔 수 없이 elif 대신 if...else 두 개로 분리하여야 한다. 다만 이 경우 가독성을 해치기 때문에 다른 방법을 활용하는 것을 추천한다.

 

list1 = [1, 2, 3, 4, 5]
new_list = ["one" if v == 1 else "two" if v == 2 else "other" for v in list1]
print(new_list)
# ['one', 'two', 'other', 'other', 'other']

 

만약 여러 조건의 if가 존재할 경우 파이썬이 암시적으로 and 로 연결하여 준다.

 

list1 = [1, 2, 3, 4, 5]
new_list1 = [x for x in list1 if x > 3 if x % 2 == 0]
new_list2 = [x for x in list1 if x > 3 and x % 2 == 0]
print(new_list1, new_list2)
# [4] [4]
# 같은 값을 반환한다!

 

new_list1에 논리 연산자 혹은 연결하여 주는 구문이 없어 오류가 발생한다고 착각하는 경우가 있는데, 실제로 실행하여보면 new_list2와 같은 논리로 같은 list를 반환하는 것을 확인할 수 있다.

 

가독성을 높이려면 새로운 함수를 작성하거나, 헬퍼 함수를 작성하는 것을 추천한다.

---

파이썬 3.8로 넘어오면서 물개 이빨 모양(:=, walrus operator)가 도입되었는데, 이 연산자와 list comprehension을 동시에 사용할 수 있다.

 

다음 코드는 list에서 문자열의 길이가 3보다 큰 문자열에 대하여 길이를 담는 list를 만든다고 할 때, 다음과 같이 활용할 수 있다.

 

strings = ["wal", "rus", "operator", "in", "python", "lang"]
 
strings_length_longer_than_three = [x for string in strings if (x := len(string)) > 3]
print(strings_length_longer_than_three)
# [8, 6, 4]

 

만약 walrus operator(:=)가 익숙하지 않다면 넘어가도 좋다. walrus operator 자체가 가독성에 대하여 논란이 많은 연산자이다.

---

3. Comprehension을 다른 곳에서도 사용하기

List comprehension의 문법은 list 이외에도 set이나 dictionary 자료형에서도 가능하다. 다만 반환되는 자료형이 set과 dict라는 점에서의 차이가 존재할 뿐이다.

 

print({x : x*x for x in range(5)})
# {0: 0, 1: 1, 2: 4, 3: 9, 4: 16}
 
print(set([i for i in [1, 2, 5, 4, 1, 2]]))
# {1, 2, 4, 5}

 

다만, 다음과 같이 tuple은 comprehension이 아닌 generator로 사용되니 주의하여야 한다!

 

print((i for i in range(5)))
# <generator object <genexpr> at 0x0000029A7ADF4190>
# tuple이 아니다!

 

이로써 list 이외의 다른 곳에서도 comprehension 문법을 사용할 수 있다는 것을 알아보았다.