백준 15652번: N과 M (4) (S3) - Python 풀이

문제 설명

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

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

• 1부터 N까지 자연수 중에서 M개를 고른 수열
• 같은 수를 여러 번 골라도 된다.
• 고른 수열은 비내림차순이어야 한다.
• 길이가 K인 수열 A가 A1 ≤ A2 ≤ ... ≤ AK-1 ≤ AK를 만족하면, 비내림차순이라고 한다.

입력

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

출력

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

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

예제 입력 1

3 1

예제 출력 1

1
2
3

예제 입력 2

4 2

예제 출력 2

1 1
1 2
1 3
1 4
2 2
2 3
2 4
3 3
3 4
4 4

예제 입력 3

3 3

예제 출력 3

1 1 1
1 1 2
1 1 3
1 2 2
1 2 3
1 3 3
2 2 2
2 2 3
2 3 3
3 3 3

풀이 과정

(2)번 문제(= 1 4랑 4 1을 같은 경우로 취급해서 4 1은 출력 X)와 (3)번 문제(= 같은 수를 여러 번 골라도 됨)를 합치면 될 것 같은데? 싶었습니다. 그래서 아래처럼 작성해서 제출했는데

N, M = map(int, input().split())
numbers = []

def backtracking():
    if len(numbers) == M:
        # (2)번 문제(= 1 4랑 4 1을 같은 경우로 취급해서 4 1은 출력 X) -> 오름차순 결과랑 같은 경우만 출력
        if sorted(numbers) == numbers:
            print(*numbers)
        return
    for i in range(1, N+1):
        # (3)번 문제(= 같은 수를 여러 번 골라도 됨) -> if 조건 제거
        numbers.append(i)
        backtracking()
        numbers.pop()

backtracking()

시간초과가 떴습니다.

하긴 for문 내에서 if 조건을 제거함으로써 예를 들어 입력이 4 2 일 때,

1 1
1 2
1 3
1 4
2 1
2 2
2 3
2 4
3 1
3 2
3 3
3 4
4 1
4 2
4 3
4 4

이 결과를 모두 살펴보게 되는데, (3) 문제는 이게 정답이라 상관없었지만 이 문제에서는 2로 시작하는 수열에서는 1이 쓰이지 않고, 3으로 시작하는 수열에서는 1,2가 쓰이지 않고, 4로 시작하는 수열에서는 1,2,3이 쓰이지 않기 때문에 이렇게 다 살펴볼 필요가 없습니다.

그래서 for문의 range 시작 값을 1로 고정하지 않고, 재귀할 때마다 값을 변경했습니다. backtracking 함수에 k 파라미터를 전달해 k부터 N까지 for문을 돌게 구현했습니다.

정답 코드

N, M = map(int, input().split())
numbers = []

def backtracking(k):
    # M개 다 찾았으면 출력 후 함수 종료
    if len(numbers) == M:
        print(*numbers)
        return
    # k부터 N까지 숫자 확인
    for i in range(k, N+1):
        # 수열에 추가
        numbers.append(i)
        # 재귀 호출 시 i를 전달
        backtracking(i)
        # 맨 뒷자리 수를 제거함으로써 전 단계(= 한 칸 전의 자릿수)로 돌아간다
        numbers.pop()

backtracking(1)

재귀 호출 시 숫자 넘길 때 i + 1이 아닌 i를 그대로 전달해야 합니다.

backtracking(i+1)로 적는다면 예를 들어 입력이 4 2 일 때

1 2
1 3
1 4
2 3
2 4
3 4

이렇게 나오게 됩니다. 이 문제에서는 1 1 / 2 2 / 3 3 / 4 4 처럼 같은 수도 추가해야 하기 때문에 i를 그대로 전달합니다.

이상입니다.