[백준 15649] N과 M (1) - 파이썬 with 백트래킹

1. 문제

https://www.acmicpc.net/problem/15649

2. 풀이

백트래킹 알고리즘의 기본 예제이다. 백트래킹은 DFS를 사용하여 풀 수 있는데, 일반적인 DFS와의 차이점은 가지치기를 한다는 점이다.

모든 경우의 수를 탐색하는 대신 조건을 걸어서 유망(promising)하지 않은 경우에는 탐색을 중지하고 이전 노드로 돌아가서 다른 경우를 탐색한다.

이 문제를 보면 m으로 자리 수를 정해놓고 각 자리에 들어가는 숫자는 중복되지 않는다. 예를 들면, n = 4, m = 2 인 경우를 보자. 자리 수는 두 개이고, 첫 자리에 1이 오는 경우에 둘째 자리에는 1을 제외한 나머지 숫자만 올 수 있다.
=> ((1, 2), (1, 3), (1, 4)) => (1, 1)은 불가능 하다.

자리 수가 두 개인 경우는 이중 반복문을 통해서 간단하게 풀 수 있지만, m이 커질 수록 반복문을 계속 중첩시킬 수는 없기 때문에 백트래킹을 사용해야 한다. 조건은 간단하다. 이미 사용(방문)한 경우라면 제외시키면 된다.

for i in range(1, n+1):
        if visited[i]:
            continue

이 개념을 가지고 코드 풀이를 시작해보자.
입력값을 받고 리스트를 만들어 둔다. s는 출력되기 위해 숫자가 쌓일 스택과 같다고 생각하면 좋다. DFS함수 풀이를 보면 이해가 될 것이다.

n, m = map(int, input().split())
s = [] 
visited = [False] * (n+1)

이제 DFS 함수를 만들어 준다. DFS는 재귀함수로 반복될 것이기 때문에 함수 출력 조건을 먼저 걸어준다. 리스트 s안에 m개의 요소가 쌓인다면 출력해주도록 한다.
=> m = 2 일 경우 s = [1, 2] 가 되면 출력한다.

def dfs():
    if len(s) == m:
        print(' '.join(map(str, s)))
        return

그 다음 위의 조건을 만족하지 않았을 경우 진행될 코드를 작성한다. 백트래킹의 조건으로 위에 설명한 방문했을 때는 제외시키도록 조건문을 걸어주어 아직 방문하지 않은 경우에만 방문해 주도록 한다. 방문하는 경우에는 방문처리 후, s에 값을 추가시키고 DFS함수를 다시 실행함으로써 다음 자리수로 넘어간다.

    for i in range(1, n+1):
        if visited[i]:
            continue
        visited[i] = True
        s.append(i)
        dfs()
        s.pop()
        visited[i] = False

n = 4, m =2 인 경우를 보자.

1. 함수를 처음 실행하면, s는 비어있기 때문에 len(s) = 0 != m 이기 때문에 첫 조건문은 실행되지 않는다.
2. for 반복문으로 들어가서 i = 1이 선택되고 visited[i] = False 이기 때문에 조건문이 또 실행되지 않는다.
3. visited[1] = True가 되고, s = [1] 이 된다.
4. 함수 재실행
5. len(s) = 1 != m 이라 조건문 패스. i = 1 부터 반복되는데 visited[1] = True이기 때문에 조건문에 걸려 continue 된다. i = 2일 경우 visited[2] = False 때문에 조건문에 걸리지 않는다. s = [1, 2], visited[2] = True 로 변경된다.
6. 함수 재실행
7. len(s) = 2 == m 이기 때문에 첫 조건문에 걸리고 '1 2' 가 출력되게 된다.
8. s에서 2가 빠지고 visited[2] = False로 초기화 된다. 이후 반복 진행.

s에서 숫자 입출과정을 괄호로 표현하자면 (1 (2) (3) (4)), (2 (3) (4) (1)) ... 라고 할 수 있다. 괄호가 열릴때( '(' ) append, 닫힐때( ')' ) pop된다고 생각하면 쉽다.

3. 코드

def dfs():
    if len(s) == m:
        print(' '.join(map(str, s)))
        return
    for i in range(1, n+1):
        if visited[i]:
            continue
        visited[i] = True
        s.append(i)
        dfs()
        s.pop()
        print(s)
        print(visited)
        visited[i] = False
            

n, m = map(int, input().split())
s = []
visited = [False] * (n+1)

dfs()

4. 다른 풀이

사실 이 문제를 처음 풀었을 때는 itertools의 permutations 함수를 사용해서 풀었다. Permutations는 배열에서 원하는 길이에 맞는 모든 조합을 구하는 함수이다. 그런데 이 풀이는 파이썬에서만 가능한 풀이이고 문제 출제 의도에 맞지 않기 때문에 백트래킹 풀이를 고려하게 되었다.

import itertools

n, m = map(int, input().split())
nums = [i for i in range(1, n+1)]

array = itertools.permutations(nums, m)

for i in array:
    for j in i:
        print(j, end = ' ')
    print()

참고로 permutations 와 비슷한 combinations 함수가 있는데, 이 둘의 차이는 중복의 유무이다. 예를들면 Combinations 함수는 (1,2), (2,1)를 같은 요소로 본다.

array = [1, 2, 3]
permutations(array, 2) => [1,2], [1,3], [2,1], [2, 3], [3, 1], [3, 2]
combinations(array, 2) => [1,2], [1,3], [2,3]