SWEA 2817번: 부분 수열의 합 (D3) - Python 풀이

 

문제 설명

SW Expert Academy

A1, A2, ... , AN의 N개의 자연수가 주어졌을 때, 최소 1개 이상의 수를 선택하여 그 합이 K가 되는 경우의 수를 구하는 프로그램을 작성하시오.

[입력]

첫 번째 줄에 테스트 케이스의 수 T가 주어진다.
각 테스트 케이스의 첫 번째 줄에는 2개의 자연수 N(1 ≤ N ≤ 20)과 K(1 ≤ K ≤ 1000)가 주어진다.
두 번째 줄에는 N개의 자연수 수열 A가 주어진다. 수열의 원소인 N개의 자연수는 공백을 사이에 두고 주어지며, 1 이상 100 이하임이 보장된다.

[출력]

각 테스트 케이스마다 ‘#x ’(x는 테스트케이스 번호를 의미하며 1부터 시작한다)를 출력하고, 부분 수열의 합이 K가 되는 경우의 수를 출력한다.

입력

1
4 3
1 2 1 2

sample_input.txt

출력

#1 4

sample_output.txt

풀이 과정

N이 20 이하로 작고, 수열에서 일부를 골라 합이 K를 만족하는 경우(조합)가 몇 개인지 세야 합니다. 그렇다면 모든 경우를 탐색해서 어떤 경우에 합이 K를 만족하는지 확인하면 되겠습니다. → 완전탐색 기반 알고리즘 생각 가능.

N이 작으니까 재귀 방식이 가능할 것 같습니다. 수열의 각 원소에게는 1) 이 원소를 선택하여 총합에 포함한다 2) 이 원소를 선택하지 않는다 2가지의 선택지가 있습니다. 각 선택지를 재귀적으로 탐색합니다.

탐색 과정 중 총합이 K를 넘어가게 된다면, 이 경우는 어차피 답에 포함되지 않기 때문에 더이상 탐색하지 않도록 합니다. (= 백트래킹의 가지치기)

정답 코드

N, K = map(int, input().split())
A = list(map(int, input().split()))
answer = 0

def DFS(index, now_sum):
    global answer
    # 가지치기: 현재 합이 K를 초과하면 더이상 볼 필요 없으니까 종료
    if now_sum > K:
        return
    # 수열의 모든 원소를 다 봤으면
    if index == N:
        # 현재 합이 K라면
        if now_sum == K:
            # 카운트 누적
            answer += 1
        # 종료
        return

    # 수열에서 현재 index 위치의 값을 선택하지 않았을 경우
    # 인덱스만 다음 칸으로 1 증가시켜서 재귀 호출
    # 현재 값을 선택하지 않았으니 총합은 그대로
    DFS(index + 1, now_sum)
    # 수열에서 현재 index 위치의 값을 선택했을 경우
    # 인덱스 다음 칸으로 1 증가시켜서 재귀 호출
    # now_sum에 수열 index 위치의 값을 더한 값으로 재귀 호출
    DFS(index + 1, now_sum + A[index])

# 0번째 인덱스에서 총합 0으로 시작
DFS(0, 0)

print(answer)

이상입니다.