리스트 컴프리헨션(List Comprehension)

이번 글에서는 파이썬에서 리스트를 간결하고 직관적으로 생성할 수 있는 문법인 리스트 컴프리헨션에 대해 작성해보려고 합니다.

 

1. 리스트 컴프리헨션

리스트 컴프리헨션은 직관적으로 리스트를 생성하는 방법입니다.

대괄호 [ ]로 감싸고 내부에 for 문과 if 문을 사용하여 반복하며 조건에 만족하는 것만 리스트로 생성할 수 있습니다.

리스트 컴프리헨션이 코딩테스트에서 많이 쓰이는 이유

1. 한 줄로 간결하게 표현
2. 가능가독성이 뛰어남
3. 시간/메모리 효율성
4. 실전 문제에서 자주 등장하는 패턴과 잘 맞음

2. 기초 사용법

0 ~ 4 사이의 정수로 리스트를 생성한다고 했을 때 리스트 컴프리헨션을 사용하지 않고 생성하는 방법입니다.

 

코드1 (리스트 컴프리헨션 사용 X)

li=[]
for i in range(5):
    li.append(i)

 

코드2 (리스트 컴프리헨션 사용 X)

list(range(5))

 

결과

[0, 1, 2, 3, 4]

---

코드 (리스트 컴프리헨션 사용 O)

[i for i in range(5)]

 

결과

[0, 1, 2, 3, 4]

 

위 코드를 보면 물론 list(range(5)) 코드가 더 효율적인 방법입니다.

하지만 리스트 컴프리헨션을 사용하는 이유는 대괄호 내에서 수식이나 함수 적용뿐만 아니라

if - else 문으로 조건을 자유롭게 줄 수 있다는 장점이 있습니다.

 

3. 응용

1) 수식이나 함수 적용

 

① 수식 적용 

 

코드

[i * 10 for i in range(5)]

 

결과 

[0, 10, 20, 30, 40]

 

변수 i가 0 ~ 4까지 반복하며 10을 곱해, 0부터  40까지 10 단위로 출력할 수도 있습니다.

 

② 함수 적용

 

코드

def test(x):
    x = str(x)+'ab'
    return x
 
[test(i) for i in range(5)]

 

결과

['0ab', '1ab', '2ab', '3ab', '4ab']

 

test 함수 : x를 입력받아 문자열로 바꾼 뒤 ab를 추가하여 반환해주는 함수입니다.

[test(i) for i in range(5)] : 0과 4 사이의 수를 test함수에 넣고 반환한 값을 리스트로 생성합니다.

 

위처럼, 리스트 컴프리헨션에서는 리스트로 만들고자 하는 값에 함수도 적용할 수 있습니다.

(수식이나 함수뿐만 아니라 변환되는 모든 것을 다 적용할 수 있습니다.)

 

2) 조건문 사용

리스트 컴프리헨션을 쓰다보면 조건문을 쓰는게 매우 편리합니다.

조건문은 리스트 컴프리헨션 왼쪽과 오른쪽 둘 다 사용할 수 있습니다.

• 오른쪽: if 문만 사용
• 왼쪽: if - else문 사용

① 오른쪽에 if문 사용

코드

[i for i in range(5) if i%2==0]

 

결과

[0, 2, 4]

 

i for i in range(5)에서 i가 0~4를 반복합니다.

이때, if i%2==0은 i가 2로 나누어 나머지가 0일 때 즉, i가 2로 나누어 떨어질 때라는 조건을 사용하여 0~4중에서 2로 나누어 떨어지는 [0, 2, 4]만 출력하도록 합니다.

 

오른쪽에 if문을 여러 개 사용할 수 있습니다.

 

코드

[i for i in range(5) if i%2==0 if i%4==0]

 

결과

[0, 4]

 

② 왼쪽에 if문 사용(else 같이 사용)

왼쪽에 if문을 사용할 때는 반드시 else와 같이 사용해야 합니다.

 

코드

[i if i%2==0 else 'odd' for i in range(5)]

 

결과

[0,'odd', 2, 'odd', 4]

 

왼쪽에 i if i%2==0 else 'odd'는 i가 2로 나누어 떨어질 때는 i를 출력하고, 그렇지 않을 때는 'odd'를 출력하라는 의미입니다. 즉, 2로 나누어 떨어지는 0, 2, 4에 대해선 그대로 출력하고, 2로 나누어 떨어지지 않는 1과 3에 대해 선 'odd'를 출력하게 합니다.

 

왼쪽에 if else 문은 여러 개를 사용할 수 있습니다. 그렇게 되면 elif와 같은 효과를 주게 됩니다.

 

코드

[i if i%2==0 else 'odd-1' if i==1 else 'odd-3'  for i in range(5)]

 

결과

[0, 'odd-1', 2, 'odd-3', 4]

 

i if i%2==0 : i가 2로 나누어 떨어지면 i를 그대로 출력합니다.

 

else 'odd-1' if i==1 : 2로 나누어 떨어지지 않으면, i가 1일 때 'odd-1' 출력합니다.

 

else 'odd-3' : 위의 두 조건이 모두 성립하지 않으면 'odd-3' 출력합니다.

 

정리하자면, i가 2로 나누어 떨어지면 그대로 출력하고, 1일 때는 'odd-1'을 출력하고 그렇지 않은 경우는 3밖에 남지 않으므로, 3일 때 'odd-3'이 출력됩니다.

 

3) for문 2개 사용

리스트 컴프리헨션을 이용하면 for문을 2개 사용함으로써 중첩 for문의 효과를 줄 수 있습니다.

 

코드

[(i,j) for i in range(2) for j in range(3)]

 

결과

[(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)]

 

i가 0,1을 반복하며, 각각 j가 0~2를 반복하게 됩니다. 즉, i가 2번 반복되는 동안, j가 3번 반복되어 총 6가지가 출력됩니다.

출력되는 i, j에 대해서 괄호로 묶어야 정상 동작합니다.

 

위의 리스트 컴프리헨션을 풀어서 작성해보면 아래와 같은 코드가 됩니다.

 

코

li=[]
for i in range(2):
    for j in range(3):
        li.append((i,j))

 

결과

[(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)]

 

리스트 컴프리헨션은 대괄호로 둘러쌓고 안에 코드를 입력하게 됩니다.

그런데 대괄호 대신에 중괄호 "{" "}"로 둘러쌓게 되면 집합 컴프리헨션(Set Comprehension)가 됩니다.

 

코드

{i for i in range(5)}

 

결과

{1, 2, 3, 4, 5}

 

이때, 딕셔너리의 형식에 맞게 키(key)와 값(value)을 콜론(":")을 사용하면 딕셔너리 컴프리헨션(Dictionary Comprehension)이 됩니다.

 

코드

{i:i for i in range(5)}

 

결과

{0: 0, 1: 1, 2: 2, 3: 3, 4: 4}

 

또한, 소괄호 "(" ")"를 사용하면 Generator expression이 됩니다. Generator의 개념을 알면 유용하게 사용할 수 있습니다.