[python] list comprehension 사용법

정현섭·2021년 4월 16일
0

이번 학기 인공지능 관련 수업을 듣게되어 파이썬을 사용하는데,
파이썬을 파이썬답게 쓰지 못하는 것 같아서 앞으로 몇 가지 pythonic한 문법들을 정리하려고 한다.

List Comprehension 이란

  • list comprehension은 어떤 list A에서 다른 list A'를 만들 때 사용하는 문법이다.
  • 기존의 C/C++/Java 처럼 빈 list를 만들고 list A를 순회하여 list A'를 만들어나가는 방식대신 사용할 수 있다.

규칙

  • [expresion for variable in iterable]
  • [expresion for variable in iterable if condition]

사용예시

예시 1 : 심플한 사용

  • 그냥 for문을 쓴 경우
squares = []
for x in range(1, 11):
    squares.append(x**2)
  • list comprehension을 쓴 경우
numbers = [x ** 2 for x in range(1, 11)]

예시 2 : if문을 사용한 경우

  • 그냥 for문을 쓴 경우
squares = []
for x in range(1, 11):
    if x % 2:
        squares.append(x*x)
  • list comprehension을 쓴 경우
numbers = [x*x for x in range(1, 11) if x % 2]

이 경우 반복 과정에서 if문에 부합되는 경우에만 expression 부분이 실행되어 list에 데이터가 추가된다. (if문이 있다고 전체 시간복잡도 가 줄어드는 것은 아님.)

예시 3 : if문 중첩

  • 그냥 for문을 쓴 경우
arr = []
for x in range(1, 11):
    if x % 2 == 0:
    	if x % 3 == 0:
        	arr.append(x)
  • list comprehension을 쓴 경우
arr = [x for x in range(1, 11) if x % 2 == 0 if x % 3 == 0]

위와 같이 if문을 여러번 중첩하여 사용할 수 도 있다.

예시 4 : if / else

  • 그냥 for문을 쓴 경우
info = []
for i in range(10):
	if i % 2 == 0:
    	info.append("Even")
    else
    	info.append("Odd")
  • list comprehension을 쓴 경우
info = ["Even" if i%2==0 else "Odd" for i in range(10)]

if / else 구문을 쓸 경우 조건문이 for문 앞으로 나오게 되는것에 주의하자.

예시 5 : 다중 for문

matrix = [
    [1, 2, 3, 4],
    [5, 6, 7, 8],
    [9, 10, 11, 12],
]

이러한 이중 list를 그냥 1차원 list로 만들려고 한다.

  • 그냥 for문을 쓴 경우
flattened = []
for row in matrix:
    for x in row:
        flattened.append(x)
  • list comprehension을 쓴 경우
flattened = [x for row in matrix for x in row]

이때 outter for문이 먼저 나온다는 것을 유의하자.

하나 더

아래 예시는 세 변의 길이가 n보다 작은 모든 직각삼각형의 세 변 길이를 구하기이다. (정수 domain에서!)

  • 그냥 for문을 쓴 경우
def pythagorean(n):
    ret = []
    for a in range(1, n + 1):
        for b in range(a, n + 1):
            for c in range(b, n + 1):
                if a * a + b * b == c * c:
                    ret.append((a, b, c))
    return ret
  • list comprehension을 쓴 경우
def pythagorean(n):
    return [(a, b, c) for a in range(1, n + 1) for b in range(a, n + 1) for c in range(b, n + 1) if a*a + b*b == c*c]

추가 : Dictionary Comprehension

fruit = ['apple', 'banana', 'kiwi', 'mango']
count = [2, 8, 1, 9]

이러한 list들이 있을 때, fruit를 key로 count를 value로 하는 dictionary를 만들려고 한다.

  • 그냥 for문을 쓴 경우
alist = []
for f, c in zip(fruit, count):
    alist.append((f, c))
store = dict(alist)
  • list comprehension을 쓴 경우
store = dict([(f, c) for f, c in zip(fruit, count)])
  • dictionary comprehension을 쓴 경우
store = {fruit[i]:count[i] for i in range(len(fruit))}

사실 위 예시처럼 iterable한 data type이면 comprehension이 가능하다. (ex. list, set, dictionary)

성능 차이

위 예시들에서 봤다시피, list comprehension 방식을 쓰면 훨신 코드를 간결하게 만들 수 있다.
그런데 코드 성능면에서도 list comprehension 방식이 더 빠르다.

결론

  • list comprehension을 쓰면 코드도 짧고 성능도 좋다.
  • 그래도 가독성은 기존 for loop이 더 좋으니, 코드를 짤 때 상황에 맞게 기존 for loop 과 list comprehension을 잘 섞어 사용하면 될 것 같다.

0개의 댓글