[python]리스트 컴프리핸션에서 if문 사용하기

python

리스트 컴프리핸션(list comprehension)이란 pythonic 하다고 표현되는 대표적인 문법 중에 하나이다. 꽤 간단한 방법으로 리스트를 만들 수 있다는 장점이 있다. 기본적으로 for문을 사용하기 때문에 반복할 수 있는 iterator나 generator와 함께 사용되는 것이 일반적이다.

list comprehension을 사용할때 if문을 사용하면 조건에 맞는 데이터를 가지고 리스트를 만들 수 있는데 이때 if문을 사용하는 방법에는 두 가지 정도가 있는 것 같다. 이 두 가지 방법에 대해서 정리한다.

조건에 맞는 리스트 만들기

만약 a라는 리스트에서 짝수인 데이터만 가지고 새로운 리스트(b)를 만들어야 한다면 아래와 같이 할 수 있다.

1
2
3
4
a = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

b = [num for num in a if num % 2 == 0]
# [2, 4, 6, 8, 10]

for문 뒤에 if문을 적어주면 된다.

삼항연상 이용

if ... else ... 삼항연산을 이용하면 조건에 해당할 때와 해당하지 않을때에 대해서 리스트 컴프리핸션을 사용할 수 있다. 짝수면 그대로 두고 홀 수 일때 10을 곱하는 컴프리핸션은 다음과 같이 작성할 수 있다.

1
2
3
4
5
6
7
a = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

b = [
    num if num % 2 == 0 else num * 10
    for num in a
]
# [10, 2, 30, 4, 50, 6, 70, 8, 90, 10]

정리

리스트를 만들때 조건이 복잡하지 않다면 리스트 컴프리핸션을 사용하는 것이 권장된다. 간단한 조건이라면 삼항연산과 함께 리스트 컴프리핸션을 사용해도 되지만 그렇지 않을 경우 조건을 판단하는 함수를 따로 만들어서 삼항 연산을 대체하거나 다른 방법을 찾는 것이 좋을 것 같다.

Share