Python: 파이썬 리스트 컴프리헨션: 2차원 리스트에서 모든 요소가 같은지 판단하기, 2차원 리스트 나누기, 2차원 리스트의 sum(총합), 2차원 리스트끼리 더하기

📖 리스트 컴프리헨션[List Comprehension] 이란?

파이썬에서 지원하는 리스트 사용 방법으로, 리스트 내에서 조건문이나 반복문을 사용하여 간결하게 리스트를 만들 수 있다.

예를 들면 숫자 1부터 9를 요소로 가지는 리스트를 만들때, 아래와 같이 사용할 수 있다.

li = [i for i in range(1, 10)] # [1, 2, 3, 4, 5, 6, 7, 8, 9]

 

다음은 2차원 리스트를 생성하는 코드이다.

li_2 = [
        [i+j for j in range(1, 4)]
        for i in range(0, 9, 3)
        ]
print(li_2) # [[1, 2, 3], [4, 5, 6], [7, 8, 9]]

위의 코드를 설명하자면 다음과 같다.

1. `for i in range(0, 9, 3)` 이 먼저 실행된다. `i`의 값은 순서대로 `0, 3, 6`이다.
2. `[i+j for j in range(1, 4)]`가 두 번째로 실행된다. `i`의 값의 수 만큼, 즉 3번 실행된다.
   • `i = 0`일 때, `[0+1, 0+2, 0+3]` ➡️ [1, 2, 3]
   • `i=3`일 때, ` [3+1, 3+2, 3+3]` ➡️ [4, 5, 6]
   • `i=6`일 때, ` [6+1, 6+2, 6+3]` ➡️ [7, 8, 9]
3. 두 번째로 실행된 리스트들이 `li_2` 에 담겨 2차원 리스트로 반환된다.

 

📌  2차원 리스트에서 리스트 컴프리헨션 사용

그렇다면, 2차원 리스트를 컴프리헨션을 이용해 편하게 다룰 수 있는 방법들을 알아보자. 2차원 리스트에서 리스트 컴프리헨션을 사용하려면, `for`문을 두 번 사용하면 된다. 리스트 컴프리헨션에서 `for` 반복문을 여러 번 사용하면, 첫 번째 반복문부터 순서대로 중첩하여 실행한다. 

1️⃣ 모든 요소가 동일한지 판단하기

2차원 리스트에서 모든 요소가 동일한지 판단하기 위해서 보통 2중 반복문을 통해 아래와 같이 코드를 짠다.

def is_same(data):
    val = data[0][0]
    for d in data:
    	for n in d:
        	if n != val: return False
    return True

이를 `all()` 함수와 리스트 컴프리헨션으로 간단하게 짤 수 있다.

def is_same(data):
    val = data[0][0]
    return all(n == val for row in data for n in row)

위의 코드를 설명하자면 아래와 같다.

1. 리스트의 첫 번째 요소를 `val` 변수에 담고, 모든 요소가 `val`과 같은 값을 가진다면 `True`를, 아니라면 `False`를 반환한다.
2. `all()` 리스트의 모든 요소가 `True`이면 `True`를, 한개라도 `False`이면 `False`를 반환하는 함수이다. 리스트 내 모든 요소에 대하여 `n`과 `val`의 값이 같은지 검증한다.
3. `for row in data`가 먼저 실행된다. 2차원 리스트인 data 에서 row를 가져온다
4. `for n in row`가 그 다음으로 실행된다. 가져온 row 안의 값(`n`)을 가져온다
5. `n`과 `val`의 값을 비교한다

 

2️⃣ 2차원 리스트를 N*N개의 블록으로 나누기

예를 들면, 크기가 4*4인 2차원 리스트를 크기가 2*2인 블록으로 나누거나, 9*9인 2차원 리스트를 3*3 블록으로 나누고 싶다면 아래와 같이 코드를 구성하게 된다.

def split_blocks(data, volum):
    n = len(data) // volum
    blocks = []
    for i in range(0, len(data), n):
    	for j in range(0, len(data), n):
        	blocks.append([row[j:j+n] for row in data[i:i+n]])
    return blocks

 

위의 코드를 리스트 컴프리헨션으로 변경하면 다음과 같다.

def split_blocks(data, volume):
    n = len(data) // volume
    return [
        [row[j:j+n] for row in data[i:i+n]]
        for i in range(0, len(data), n)
        for j  in range(0, len(data), n)
    ]

temp = [[i for i in range(9)] for _ in range(9)]
blocks = split_blocks(temp,3)

1. 크기가 9*9인 2차원 리스트 `temp`를 선언
2. `split_blocks()`를 통해 3*3 블록 9개로 나눈다.
   • `for i in range(0, len(data), n)` i가 `n` 단위로 반복
   • `for j in range(0, lend(data), n)` j가 `n` 단위로 반복
     • `[row[j:j+n] for row in data[i:i+n]]`이 i값의 범위 만큼 생성
     • i = 0, j = 0, 3, 6: [row[0:3] for row in data[0:3]], [row[3:6] for row in data[0:3]], [row[6:9] for row in data[0:3]]
     • i = 3, j = 0, 3, 6: [row[0:3] for row in data[3:6]], [row[3:6] for row in data[3:6]], [row[6:9] for row in data[3:6]]
     • i = 6, j = 0, 3, 6:  [row[0:3] for row in data[6:9]], [row[3:6] for row in data[6:9]], [row[6:9] for row in data[6:9]]

잘 이해가 되지 않는다면 아래의 그림을 보면 쉽다. 

2차원 리스트 블록 나누기

블록을 나누는 코드를 리스트 컴프리헨션으로 짜면 몇 자 줄어들지만, 헷갈릴 수 있을 것 같다. 가독성이 그렇게 좋지도 않은 것 같아서 개인적으로 잘 사용하지는 않을 것 같다. 아니다. 엄청 잘 쓴다. ㅋㅋㅋㅋㅋㅋ

 

3️⃣ 2차원 리스트의 총합 구하기

`sum()`을 이용하여 2차원 리스트의 총합을 구할 수 있다. 

li = [[1] * 3 for _ in range(3)] # [[1, 1, 1], [1, 1, 1], [1, 1, 1]]
print(sum(li, [])) # [1, 1, 1, 1, 1, 1, 1, 1, 1]
print(sum(sum(li, []))) # 9

1. `sum(li, [])`: sum에 첫 번째 인자로 2차원 리스트를 던지고, 두 번째 인자로 빈 리스트를 던지면 2차원 리스트끼리 합해서 1차원 리스트를 반환한다. 
2. `sum(sum(li, []))`: `sum(li, [])`를 활용해 1차원 리스트를 구하고, 그 합계를 구한다

개인적으로 제일 많이 쓰는 코드다...

 

4️⃣ 2차원 리스트를 각 요소끼리 더하기

길이가 같은 두 2차원 리스트가 있을때, `zip()`을 활용해서 각 요소끼리 더하고, 그 결과를 2차원 리스트로 반환할 수 있다.

arr1 = [[1, 2], [3, 4]]
arr2 = [[5, 6], [7, 8]]
sum_arr = [[a + b for a, b in zip(arr1[i], arr2[i])] for i in range(len(arr1))]
print(sum_arr) # [[6, 8], [10, 12]]

1. `for i in range(len(arr1))`이 먼저 실행된다. i = 0, 1
2. `[a + b for a, b in zip(arr1[i], arr2[i])]`이 실행된다
   • `zip()`을 통해 리스트 두개의 각 자리의 요소를 가져와서 더한다. 
   • [`1+5`, `2+6`], [`3+7`, `4+8`]
3. 결과는 2차원 리스트이다. `[[6, 8], [10, 12]]`

 

📌 정리

지금까지 리스트 컴프리헨션을 이용하여 코드를 단순화 하는 방법을 알아봤다. 위에 코드들은 코딩테스트 등 쓸 일이 많을테니 기억해두고 적재적소에 잘 쓰도록 하자. 그리고, 리스트 컴프리헨션을 과도하게 사용하면 가독성이 떨어지는 경우가 있으니 적당히 쓰자.