[Python] 백준 15664 N과 M (10) (Backtracking)

📌 문제

N개의 자연수와 자연수 M이 주어졌을 때, 아래 조건을 만족하는 길이가 M인 수열을 모두 구하는 프로그램을 작성하시오.

N개의 자연수 중에서 M개를 고른 수열
고른 수열은 비내림차순이어야 한다.
길이가 K인 수열 A가 A1 ≤ A2 ≤ ... ≤ AK-1 ≤ AK를 만족하면, 비내림차순이라고 한다.

입력

첫째 줄에 N과 M이 주어진다. (1 ≤ M ≤ N ≤ 8)

둘째 줄에 N개의 수가 주어진다. 입력으로 주어지는 수는 10,000보다 작거나 같은 자연수이다.

출력

한 줄에 하나씩 문제의 조건을 만족하는 수열을 출력한다. 중복되는 수열을 여러 번 출력하면 안되며, 각 수열은 공백으로 구분해서 출력해야 한다.

수열은 사전 순으로 증가하는 순서로 출력해야 한다.

예제 입력 1

3 1
4 4 2

예제 출력 1

2
4

예제 입력 2

4 2
9 7 9 1

예제 출력 2

1 7
1 9
7 9
9 9

예제 입력 3

4 4
1 1 2 2

예제 출력 3

1 1 2 2

---

📌 풀이

💬 Code

import sys
input = sys.stdin.readline

def dfs(idx):
    if len(arr) == m:
        print(' '.join(map(str, arr)))
        return
    former = 0
    for i in range(idx+1, n):
        if not isUsed[i] and nums[i] != former:
            arr.append(nums[i])
            isUsed[i] = 1
            former = nums[i]
            dfs(i)
            arr.pop()
            isUsed[i] = 0

n, m = map(int, input().split())
nums = sorted(list(map(int, input().split())))
arr = []
isUsed = [0] * n

dfs(-1)

💡 Solution

백트래킹 연습을 위해 itertools를 사용하지 않고 풀어보았다. 차근차근 조건들을 하나씩 추가해보자.

• nums: 숫자들을 담은 리스트
• arr: 정답 리스트
• isUsed: 중복 제거를 위한 사용여부 boolean 리스트
• former: 중복 정답 제거를 위한 이전 숫자 기억 변수

1단계

백트래킹 기본 틀 잡기

def dfs():
    if len(arr) == m:
        print(' '.join(map(str, arr)))
        return
    for i in range(n):
        arr.append(nums[i])
        dfs()
        arr.pop()

# 입력
4 2
9 7 9 1

# 출력
1 1
1 7
1 9
1 9
7 1
7 7
7 9
7 9
9 1
9 7
9 9
9 9
9 1
9 7
9 9
9 9

2단계

isUsed 불린리스트를 이용해 해당 인덱스의 원소를 사용했는지를 체크해서 사용하지 않은 인덱스의 원소만 처리한다. nums=[1, 7, 9, 9]에서 현재 0번 인덱스를 사용한 경우, 이어 나올 숫자는 1, 2, 3번 인덱스에서만 가능하도록 하는 것이다.

def dfs():
    if len(arr) == m:
        print(' '.join(map(str, arr)))
        return
    for i in range(n):
        if not isUsed[i]:
            arr.append(nums[i])
            isUsed[i] = 1
            dfs()
            arr.pop()
            isUsed[i] = 0

# 입력
4 2
9 7 9 1

# 출력
1 7
1 9
1 9
7 1
7 9
7 9
9 1
9 7
9 9
9 1
9 7
9 9

3단계

nums = [1, 7, 9, 9]처럼 같은 숫자가 여러 번 등장하는 리스트의 경우 9가 포함된 조합이 여러 번 출력되는 것을 막아야 한다. former 변수를 추가해서 바로 앞에 append된 숫자를 기억시키고 그 숫자와 다른 숫자가 등장할 경우만 처리하도록 한다.

def dfs():
    if len(arr) == m:
        print(' '.join(map(str, arr)))
        return
    former = 0
    for i in range(n):
        if not isUsed[i] and nums[i] != former:
            arr.append(nums[i])
            isUsed[i] = 1
            former = nums[i]
            dfs()
            arr.pop()
            isUsed[i] = 0

# 입력
4 2
9 7 9 1

# 출력
1 7
1 9
7 1
7 9
9 1
9 7
9 9

4단계

마지막으로, 비내림차순으로 만들어주어야 한다. dfs 함수에 idx 파라미터를 추가해서 재귀함수에 현재 인덱스를 넘겨주고 해당 재귀함수에서 현재 인덱스+1의 위치부터 탐색을 시작하도록 한다.

def dfs(idx):
    if len(arr) == m:
        print(' '.join(map(str, arr)))
        return
    former = 0
    for i in range(idx+1, n):
        if not isUsed[i] and nums[i] != former:
            arr.append(nums[i])
            isUsed[i] = 1
            former = nums[i]
            dfs(i)
            arr.pop()
            isUsed[i] = 0
            
dfs(-1)

# 입력
4 2
9 7 9 1

# 출력
1 7
1 9
7 9
9 9