[알고리즘] Python zip함수와 2차원 배열 회전에 대하여

1. zip()이란?

: 인덱스에 있는 요소들을 결합하여 새로운 iterable한 객체를 생성하는 함수이다.

 

코드로 빠르게 확인해 보자!

nums = [1, 2, 3]
strings = ["a", "b", "c"]

new_obj = zip(nums, strings)

print(type(new_obj)) #<class 'zip'>
print(list(new_obj)) #[(1, 'a'), (2, 'b'), (3, 'c')]

nums 배열과 strings 배열에서 각각의 인덱스에 맞게 요소들이 매핑되어 새로운 객체를 생성해 낸 것을 볼 수 있다.

 

단, 주의해야 할 점이 있는데 만약 서로 다른 두 길이의 객체를 매핑시키려 하면 가장 짧은 길이의 객체를 기준으로 매핑을 하게 되고 나머지는 버려진다.

nums = [1, 2, 3, 4, 5]
strings = ["a", "b", "c"]

new_obj = zip(nums, strings)

print(type(new_obj)) #<class 'zip'>
print(list(new_obj)) #[(1, 'a'), (2, 'b'), (3, 'c')]

nums 배열의 길이는 5이지만 더 짧은 strings배열의 길이에 맞게 매핑이 이루어지고 nums[3]과 nums[4]는 버려졌다.

 

 

2. zip함수와 2차원 배열 회전의 연관성

자, 그럼 이게 2차원 배열 회전과 무슨 연관이 있나?

 

백준 알고리즘 사이트의 1652번 문제이다.
https://www.acmicpc.net/problem/1652

1652번: 누울 자리를 찾아라

알고리즘 풀이 포스트가 아니기 때문에 문제에대해서 링크만 걸어놓았다.

해당 문제에 대해서 내 풀이는 아래 코드와 같았다.

import sys

input = sys.stdin.readline

N = int(input())

graph = []
for _ in range(N):
    a = list(input().rstrip())
    graph.append(a)


def count_ans(graph):
    ans = 0
    for row in graph:
        count = 0
        for cell in row:
            if cell == ".":
                count += 1
            else:
                if count >= 2:
                    ans += 1
                count = 0
        if count >= 2:
            ans += 1

    return ans


def transpose(graph):
    transposed_graph = []
    for i in range(N):
        transposed_row = []
        for j in range(N):
            transposed_row.append(graph[j][i])
        transposed_graph.append(transposed_row)

    return transposed_graph


ans_row = count_ans(graph)

transposed_graph = transpose(graph)
ans_col = count_ans(transposed_graph)

print(ans_row, ans_col, end=" ")

(bfs로도 풀 수 있었지만 그냥 쌩으로 한 번 풀어보았다.)

 

transpose함수는 NxN 그래프를 전치시키는 역할을 하는 함수인데

만약 문제의 예제대로 행과 열에 대해서 각각 "."이 연속되는 것을 구하기 위해 transpose함수를 사용하게 되면 행과 열이 완전히 뒤바뀐 것을 확인할 수 있다.

....X	  	....X
..XX.	  	...X.
.....	->	.X.X.
.XX..	  	.X...
X....	  	X....

 

코딩테스트를 보다 보면 너무 긴장해서 나처럼 아예 전치된 그래프를 구하는 함수를 구현할 수 있는데 이를 zip함수를 사용하여 아주 쉽게 바꿀 수 있다.

def transpose(graph):
    transposed_graph = list(zip(*graph))
    return transposed_graph

기존 코드에서 transpose 함수만 변경되었는데 결과는 동일하게 출력이 된다.

 

이게 어떻게 이렇게 된 것인지 코드를 파헤쳐보면

graph = [
	['.', '.', '.', '.', 'X'], 
	['.', '.', 'X', 'X', '.'], 
	['.', '.', '.', '.', '.'], 
	['.', 'X', 'X', '.', '.'], 
	['X', '.', '.', '.', '.']
]

transpose_graph = [
	('.', '.', '.', '.', 'X'), 
	('.', '.', '.', 'X', '.'),
	('.', 'X', '.', 'X', '.'), 
	('.', 'X', '.', '.', '.'), 
	('X', '.', '.', '.', '.')
]

graph는 2차원 배열이기 때문에 각 행(인덱스 요소)들 또한 iterable한 객체이고, 각 행의 같은 인덱스 요소들은 zip함수로 묶어서 행과 열이 뒤바뀐 그래프를 생성할 수 있는 것이었다.

즉, list(zip(*graph))는 graph의 모든 행에 대하여 같은 인덱스 요소를 매핑시켜 새로운 행을 만들어내는 함수이다.

(graph의 1열 -> transpose_graph의 1행,
graph의 2열 -> transpose_graph의 2행,
graph의 3열 -> transpose_graph의 3행,

...)

이렇게 zip 함수의 특성을 이용하면 행과 열을 뒤 바꿀 수 있는 그래프를 아주 쉽게 구해낼 수가 있다.

zip함수의 존재는 알고 있었지만 알고리즘 풀이에서 이렇게 쓰일 줄은 정말 생각도 못했다!
오늘도 또 하나 배워간다 :)