[Python] 2차원 배열에서의 조합

della·2022년 1월 26일

주어진 2차원 배열에서 값이 0인 데이터의 좌표만을 3개씩 묶어 중복이 없는 조합의 형태로 나타내고 싶을 때, 어떻게 코딩하는것이 좋을까?

다음과 같은 입력을 받았다고 가정하고 코딩을 시작해보자.

n, m = 2, 2  # 가로, 세로
arr = [[0, 0], [0, 1]]  # 2차원 배열

1. 단순 반복문

단순하게 생각할 수 있는 방법 중 하나이다. 각 반복문마다 좌표를 하나씩 담당하여 확인하고, 모든 좌표가 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 지옥 또한 서비스로 맛볼 수 있다....

2. DFS 1

이쯤에서 코딩 테스트를 좀 공부해 봤다면, 백트래킹을 통해 순열, 조합 등을 처리할 수 있다는 것을 알 것이다.

백트래킹의 대표적인 알고리즘 중 하나인 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 라는 부분이 의미하는 것이 무엇인지 파악하는 게 어려울 수 있다. 간단히 얘기하면 중간에 값이 생략되는 것을 막아주는 것인데, 한 스텝 한 스텝씩 확인해가며 보면 이해가 쉽다.

  1. (0, 0), (0, 1), (1, 0), (1, 1) 과 같이 총 4군데 방문할 수 있는 좌표가 있다.

  2. dfs(0, 0)가 호출되면, (0, 0)이 스택에 추가되고 dfs(0, 1)를 호출한다.

    stack = [(0, 0)], x = 0, y = 1

  3. dfs(0, 1)가 호출되면, (0, 1)이 스택에 추가되고 dfs(0, 2)를 호출한다.

    stack = [(0, 0), (0, 1)], x = 0, y = 2

  4. dfs(0, 2)가 호출되면, (0, 2)는 존재하지 않는 좌표이므로 바깥쪽 for을 돌아 x가 1 증가된다

    stack = [(0, 0), (0, 1)], x = 1, y = 2

  5. 안쪽 for문을 돌기위해 조건을 검사하니 y가 여전히 2이므로, 반복문을 돌 수 없다. sy = sy if x == sx else 0 이 조건을 통하여 sy를 0으로 다시 바꿔주고나면 정상적으로 (1, 0)에 방문할 수 있게 된다.

    stack = [(0, 0), (0, 1), (1, 0)]

  6. 나머지 과정 생략...

이제 어느 정도 이해가 되지만, 여전히 복잡해 보이는 것은 어쩔 수 없다. 다른 방법을 알아보자.

3. DFS 2

일반적인 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 값을 직접 계산하여 위에 이보다 좀 더 직관적인 코드가 되었다. 하지만 결국에 “재귀”가 들어가면 모든 코드가 복잡해 보이는 건 어쩔 수 없다. 더 간단하게 만들어보자.

4. 내장 라이브러리 사용하기

가장 간단한 방법이다. 위의 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 메서드를 통해 조합을 만들어낸다.

profile
안녕하세요

0개의 댓글