주어진 2차원 배열에서 값이 0인 데이터의 좌표만을 3개씩 묶어 중복이 없는 조합의 형태로 나타내고 싶을 때, 어떻게 코딩하는것이 좋을까?
다음과 같은 입력을 받았다고 가정하고 코딩을 시작해보자.
n, m = 2, 2 # 가로, 세로
arr = [[0, 0], [0, 1]] # 2차원 배열
단순하게 생각할 수 있는 방법 중 하나이다. 각 반복문마다 좌표를 하나씩 담당하여 확인하고, 모든 좌표가 0인 경우에만 append를 실시한다.
result = []
for i in range(0, n * m - 2):
x1, y1 = i // m, i % m
for j in range(i + 1, n * m - 1):
x2, y2 = j // m, j % m
for k in range(j + 1, n * m):
x3, y3 = k // m, k % m
if arr[x1][y1] == arr[x2][y2] == arr[x3][y3] == 0:
result.append([(x1, y1), (x2, y2), (x3, y3)])
print(result)
위 코드의 가장 큰 문제점은, 시간이 오래 걸리고 코드도 복잡하다는 것이다. 사실상 좋은 게 하나도 없다는 뜻...
물론 if 문을 각 for 문을 돌기 전에 따로따로 검사를 진행하면 속도는 훨씬 빨라지지만, indent 지옥 또한 서비스로 맛볼 수 있다....
이쯤에서 코딩 테스트를 좀 공부해 봤다면, 백트래킹을 통해 순열, 조합 등을 처리할 수 있다는 것을 알 것이다.
백트래킹의 대표적인 알고리즘 중 하나인 DFS를 사용하여 풀면 다음과 같은 코드가 나온다.
def dfs(sx, sy, stack=[], result=[]):
if len(stack) == 3:
result.append(stack.copy())
return
for x in range(sx, n):
sy = sy if x == sx else 0
for y in range(sy, m):
if arr[x][y] == 0:
stack.append((x, y))
dfs(x, y + 1, stack, result)
stack.pop()
return result
print(dfs(0, 0))
전체적으로 이전보다 깔끔한 코드가 나왔다. 각 좌표를 한 번씩만 방문해서 값이 0이라면 stack에 넣고, stack이 3개가 되었을 때 결과 배열에 담아주는 방식으로 작동한다.
코드를 보다 보면 sy = sy if x == sx else 0
라는 부분이 의미하는 것이 무엇인지 파악하는 게 어려울 수 있다. 간단히 얘기하면 중간에 값이 생략되는 것을 막아주는 것인데, 한 스텝 한 스텝씩 확인해가며 보면 이해가 쉽다.
(0, 0), (0, 1), (1, 0), (1, 1) 과 같이 총 4군데 방문할 수 있는 좌표가 있다.
dfs(0, 0)가 호출되면, (0, 0)이 스택에 추가되고 dfs(0, 1)를 호출한다.
stack = [(0, 0)], x = 0, y = 1
dfs(0, 1)가 호출되면, (0, 1)이 스택에 추가되고 dfs(0, 2)를 호출한다.
stack = [(0, 0), (0, 1)], x = 0, y = 2
dfs(0, 2)가 호출되면, (0, 2)는 존재하지 않는 좌표이므로 바깥쪽 for을 돌아 x가 1 증가된다
stack = [(0, 0), (0, 1)], x = 1, y = 2
안쪽 for문을 돌기위해 조건을 검사하니 y가 여전히 2이므로, 반복문을 돌 수 없다. sy = sy if x == sx else 0
이 조건을 통하여 sy를 0으로 다시 바꿔주고나면 정상적으로 (1, 0)에 방문할 수 있게 된다.
stack = [(0, 0), (0, 1), (1, 0)]
나머지 과정 생략...
이제 어느 정도 이해가 되지만, 여전히 복잡해 보이는 것은 어쩔 수 없다. 다른 방법을 알아보자.
일반적인 1차원 배열에서의 조합 방식을 그대로 채택하여 사용하기 위해 for 문을 하나로 줄이고 x, y를 현재 인덱스 값으로 계산할 수 있도록 변경하였다. 인자도 depth 하나만 받는다.
def dfs(depth, stack=[], result=[]):
if len(stack) == 3:
result.append(stack.copy())
return
for i in range(depth, n * m):
x, y = i // m, i % m
if arr[x][y] == 0:
stack.append((x, y))
dfs(i + 1, stack, result)
stack.pop()
return result
print(dfs(0))
1번의 단순 반복 문과같이, for 문을 (n X m) 만큼 한 번에 돌면서, x 와 y 값을 직접 계산하여 위에 이보다 좀 더 직관적인 코드가 되었다. 하지만 결국에 “재귀”가 들어가면 모든 코드가 복잡해 보이는 건 어쩔 수 없다. 더 간단하게 만들어보자.
가장 간단한 방법이다. 위의 dfs를 이용한 방법보다 다소 느릴 수 있으나, 결과적으로 코드가 깔끔하기 때문에 가장 선호되는 방법이라고 볼 수 있다.
from itertools import combinations
coords = [(x, y) for x in range(n) for y in range(m) if arr[x][y] == 0]
result = list(combinations(coords, 3))
print(result)
coords에 방문할 수 있고 값이 0인 좌표에 대한 정보만이 담기고, combinations 메서드를 통해 조합을 만들어낸다.