[파이썬/Python] 리스트 안에서 반복문, Comprehension

요즘 틈 나는 대로 프로그래머스의 코딩 테스트 문제를 풀고 있다. 지금 Python으로 풀 수 있는 Level 1 문제들은 세 문제만을 남겨두고 있으며, 슬슬 Level 2 문제들도 풀어나가고 있다. 코드 정리는 귀찮아서 못하고 있는데 일단 Git에만 푸쉬해두고 있는 상태다. (Git에도 사실 다 푸쉬는 못했다..ㅎ) 아 글은 언제 쓰냐아으ㅏ으아아ㅏㅏ 귀찮다.

 

아직 저레벨 수준의 코딩 테스트 문제들이었지만, 문제를 풀어나가며 유용했던 함수 및 메소드를 정리해볼까 한다. 오늘은 그 첫 번째 순서로 Python에서 아주 유용한 comprehension이다.

Comprehension

Comprehension이란, iterable한 객체를 생성하는 방법 중 하나로 아주 유용하게 쓰고 있다. 나는 list랑 dictionary를 만들 때만 사용했지만, set을 만들 때에도 사용이 가능하다고 한다. 사용 예시를 여러 개 들어보며 comprehension에 대한 이해도를 높여보자. 예시는 list로 들겠다.

(1) for문만 들어간 comprehension

만약 0이 10개 들어가 있는 리스트를 만들고 싶다고 가정할 때, 일반적인 for문을 활용하면 다음과 같은 코드를 짤 수 있다.

arr = []
for i in range(10):
	arr.append(0)
    
print(arr) # [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

하지만 comprehension을 이용하면 다음과 같이 한 줄 만에 코드를 짤 수 있다.

arr = [0 for i in range(10)]

print(arr) # [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

위 comprehension을 이용하여 아래와 같은 코드들로 응용해볼 수도 있다.

# 예시1. 1부터 10까지의 숫자로 구성된 리스트 생성
arr1 = [i for i in range(1,11)] # arr1 = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

# 예시2. 1부터 10까지의 숫자 중 짝수로만 구성된 리스트 생성
arr2 = [i for i in range(2, 11, 2)] # arr2 = [2, 4, 6, 8, 10]

# 예시3. 1부터 10까지 각 숫자의 제곱으로 구성된 리스트 생성
arr3 = [i*i for i in range(1, 11)] # arr3 = [1, 4, 9, 16, 25, 36, 49, 64, 81, 100]

# 예시4. 기존의 리스트 arr3에 10을 더한 숫자로 구성된 리스트 생성
arr4 = [10 + i for i in arr3] # arr4 = [11, 14, 19, 26, 35, 46, 59, 74, 91, 110]

 

(2) for문과 함께 if구문, 혹은 if-else구문이 들어간 comprehension

만약 기존의 리스트 [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]에서 짝수인 원소만 추출하여 새로운 리스트를 만들고 싶다고 가정할 때, 일반적인 for문을 활용하면 다음과 같은 코드를 짤 수 있다.

arr = [1,2,3,4,5,6,7,8,9,10]
lst = []
for i in arr:
	if i % 2 == 0:
		lst.append(i)
        
print(lst) # [2, 4, 6, 8, 10]

하지만 comprehension을 이용하면 이것도 마찬가지로 한 줄로 코드를 짤 수 있다.

arr = [1,2,3,4,5,6,7,8,9,10]
lst = [i for i in arr if i % 2 == 0]

print(lst) # [2, 4, 6, 8, 10]

위 코드에서 보이는 것과 같이 comprehension을 활용할 때 if 구문을 쓸 수 있다. if 구문을 활용한 예시를 몇 개 더 들어보자. 심지어 list 안에 else 구문까지 함께 쓸 수 있는데, 이때 주의할 점은 if-else구문이 for문보다 앞에 와야 한다는 점이다.

# 예시1. 기존의 리스트에서 int형 원소만 추출하여 새로운 리스트 생성
arr = ['a', 'b', 'c', 1, 3, 12, 'd']
lst = [i for i in arr if type(i) == int] # lst = [1, 3, 12]

# 예시2. 
# 기존의 리스트에 숫자가 담겨 있지만 간혹 어떤 숫자들은 str형으로 담겨 있다.
# 기존의 리스트는 변형시키지 않고 새로운 리스트의 원소를 모두 int형으로 바꿔보자.

arr = [1, '2', '3', 4, 5]
lst = [int(i) if type(i) == str else i for i in arr] # lst = [1, 2, 3, 4, 5]

if 구문을 이용할 때 알아둘 점이 있는데, if 구문을 두 개 이상 쓰면 각 if 뒤에 있는 조건들은 모두 and로 취급되어 모든 조건들을 만족시켜야 한다. 나는 굳이 if를 여러 개 써서 가독성을 나쁘게 만들 바에야 and로 처리하는 것이 맞다고 생각한다.

arr = [i for i in range(16)] # arr = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]

# 아래 두 리스트는 동일
lst1 = [i for i in arr if i % 2 == 0 and i % 3 == 0] # lst1 = [0, 6, 12]
lst2 = [i for i in arr if i % 2 == 0 if i % 3 == 0] # lst2 = [0, 6, 12]

 

(3) for문이 2개 이상 들어간 comprehension

comprehension에서 for문을 2개 이상 사용할 수도 있다. 개인적으로는 for문을 2개까지만 쓰는 것이 좋다고 생각하며, 3개가 넘어가는 for문을 이용해야 될 경우에는 comprehension보다는 일반적인 for문을 사용하는 것이 자신한테도, 같이 개발하는 동료들한테도 좋다고 생각한다. 이해하기가 많이 어려워지기 때문이다. 아래 예시들을 통해 for문이 2개 이상 들어간 comprehension을 알아보자.

 

만약 [[1, 2, 3], [4, 5, 6], [7, 8, 9]]로 구성되어 있는 이중 리스트에서 모든 숫자를 하나의 새로운 리스트의 원소로 넣고 싶다고 가정해보자. 일반적인 이중 for문을 활용하면 다음과 같이 진행될 것이다.

arr = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
lst = []
for i in arr:
	for j in i:
		lst.append(j)
		
print(lst) # [1, 2, 3, 4, 5, 6, 7, 8, 9]

위 코드와 똑같은 기능을 수행하게끔 comprehension을 이용하면 다음과 같이 코드를 짤 수 있다.

lst = [j for i in arr for j in i]

print(lst) # [1, 2, 3, 4, 5, 6, 7, 8, 9]

위 코드에서 알 수 있듯이 먼저 나온 for이 더 높은 레벨의 for문이다. 즉, 'for i in arr'가 첫 번째 for문이고, 'for j in i'가 두 번째 for문이다.

 

그렇다면 또 다른 예시를 들어보자. arr와 똑같은 형태의 새로운 이중리스트를 만들되, arr 내의 각 int형 변수들의 제곱한 결과값을 요소로 넣고 싶다면 어떻게 할까? 아래의 코드처럼 짜면 된다.

lst = [[j*j for j in i] for i in arr]

print(lst) # [[1, 4, 9], [16, 25, 36], [49, 64, 81]]

여기서 중요한 것은, 더 나중에 수행되어야 할 for문이 먼저 나왔다는 것이다. 위 코드를 세부적으로 뜯어보면 다음과 같다.

 

1. 더 먼저 수행되어야 할 for i in arr 가 수행되면서 리스트 'arr'의 0번 원소 'arr[0]', 1번 원소 'arr[1]', 2번 원소 'arr[2]'를 순차적으로 i에 담게 된다. 이렇게 더 먼저 수행된 for문이 한 번 돌 때, 이중 리스트를 만들기 위해 [ ] 로 감싸진 'j*j for j in i'가 수행된다.

 

2. 만약 첫 번째 for문이 처음 돌 때는 i = arr[0] = [1, 2, 3] 이므로 두 번째 for문인 for j in i 에서 j는 1, 2, 3을 순차적으로 돌게 된다. 그렇게 순차적으로 돌면서 j*j값을 [ ] 안에 담게 되는데, 그 값은 1*1, 2*2, 3*3 이 될 것이다.

 

3. 그렇게 for j in i 가 모두 끝나게 되면 [1, 4, 9]라는 리스트가 새롭게 생성될 것이고, 그 리스트는 우리가 만들고자 하는 리스트 'lst'의 첫 번째 원소로 append된다.

 

4. 이후 첫 번째 for문 for i in arr 이 한 번 더 수행하게 되면서 i = arr[1] = [4, 5, 6] 이 되고, 그 이후는 2~3번의 과정을 반복하게 된다.

 

5. 결과적으로 새롭게 만든 리스트 'lst'에는 [1, 4, 9], [16, 25, 36], [49, 64, 81] 라는 세 개의 리스트가 원소로 들어가 있게 된다.

 

 

또 다른 예시로는, 조합을 만들어 볼 수도 있다. 사실 조합은 python 내장 라이브러리인 itertools를 활용하면 더 편하게 만들어낼 수 있다. 하지만 comprehension이 이렇게 유용하게 쓰일 수도 있다는 예시를 들기 위한 것이니 일단은 만들어보자.

arr1 = ['1', '3']
arr2 = ['5', '2', '8']
arr3 = ['9', '8', '5', '8']
lst = [[a, b, c] for a in arr1 for b in arr2 for c in arr3] # a,b,c를 튜플로 담아도 됨
print(lst)
# [['1', '5', '9'], ['1', '5', '8'], ['1', '5', '5'], ['1', '5', '8'], 
#  ['1', '2', '9'], ['1', '2', '8'], ['1', '2', '5'], ['1', '2', '8'], 
#  ['1', '8', '9'], ['1', '8', '8'], ['1', '8', '5'], ['1', '8', '8'], 
#  ['3', '5', '9'], ['3', '5', '8'], ['3', '5', '5'], ['3', '5', '8'], 
#  ['3', '2', '9'], ['3', '2', '8'], ['3', '2', '5'], ['3', '2', '8'], 
#  ['3', '8', '9'], ['3', '8', '8'], ['3', '8', '5'], ['3', '8', '8']]



# 필요에 따라서는 lst의 원소 형태를 아래와 같이 바꿀 수 있다!

# (1) 각 숫자를 이어 붙여서 원소로 담기
lst = [''.join([a, b, c]) for a in arr1 for b in arr2 for c in arr3]
print(lst)
# ['159', '158', '155', '158', '129', '128', '125', '128', 
#  '189', '188', '185', '188', '359', '358', '355', '358', 
#  '329', '328', '325', '328', '389', '388', '385', '388']

# (2) 각 숫자를 이어 붙인 후 int형 원소로 담기
lst = [int(''.join([a, b, c])) for a in arr1 for b in arr2 for c in arr3]
print(lst)
# [159, 158, 155, 158, 129, 128, 125, 128, 189, 188, 185, 188, 
#  359, 358, 355, 358, 329, 328, 325, 328, 389, 388, 385, 388]

코딩 테스트 문제를 풀다 보면 위와 같은 짓(?)들을 해야 될 때가 많이 있다. 그래서 comprehension이 상당히 유용하게 쓰였다.

 

쓰던 for문 쓰던 대로 쓰면 되지, 굳이 번거롭게 comprehension을 써야 한다고 물으신다면... 사실 그렇게 일반적인 for문으로 써도 무방하다. 하지만 이러한 방식이 편해지는 순간, 효율성이 상당히 올라갈 것이라고 말씀드리고 싶다. 그리고 python을 좀 더 python답게 쓰는, pythonic code를 짤 때는 comprehension이 빠질 수 없다고 생각한다.

 

글을 마무리하며, 아래 사진은 comprehension을 어떤 식으로 쓰면 될지 일반식(?)같이 정리해 본 것인데 사실 쓸 수 있는 형태가 너무 다양해서 저 아래 사진들로는 표현할 수 없는 것들도 있을 것이다. 그래도 최대한 많은 경우의 수를 담아낼 수 있게 정리해본 것이니 이 글을 보시는 분들에게 유용했으면 한다.

다만 2개 이상의 for문이 들어가는 형태의 comprehension은 너무나도 다양한 형태가 존재할 수 있다고 생각되어 정리하지 못했다.ㅠㅠ