백준 1208번: 부분수열의 합 2 (Python)

문제

N개의 정수로 이루어진 수열이 있을 때, 크기가 양수인 부분수열 중에서 그 수열의 원소를 다 더한 값이 S가 되는 경우의 수를 구하는 프로그램을 작성하시오

해설

발상이 어려웠던 문제다. 수열의 최대 크기가 40인데, 백트래킹(브루트포스)로 전체 경우를 구하게 되면 O(n^40)으로 시간초과가 발생한다. 이때, 수열을 절반으로 나누어 각 수열마다 O(n^20) + O(n^20) 의 시간복잡도를 가지게 하면 시간 내에 풀 수 있게 된다.

 

먼저, 다음과 같이 길이가 10인 수열이 있고, 합이 0이 되는 경우를 구한다고 가정해보자.

# Input
10 0
1 2 3 4 5 -1 -2 -3 -4 -5

 

먼저 수열을 절반으로 나눌 것이다.

그러면 다음과 같이 left, right 두 개의 리스트를 얻을 수 있다.

 

# left list
arr[0:mid] 
= [1 2 3 4 5]

# rightlist
arr[mid:end] 
= [-1 -2 -3 -4 -5]

 

그런 다음 left 리스트에서 나올 수 있는 모든 부분수열의 합을 구한다. 그러면 다음과 같이 나올 것이다.

 

[0, 1, 3, 6, 10, 15, 11, 7, 12, 8, 4, 8, 13, 9, 5, 10, 6, 2, 5, 9, 14, 10, 6, 11, 7, 3, 7, 12, 8, 4, 9, 5]

right 리스트의 모든 부분수열의 합도 마찬가지로 구하면..

[0, -1, -3, -6, -10, -15, -11, -7, -12, -8, -4, -8, -13, -9, -5, -10, -6, -2, -5, -9, -14, -10, -6, -11, -7, -3, -7, -12, -8, -4, -9, -5]

와 같이 나올 것이다.

 

이제 이렇게 나온 sum 값의 횟수를 map 형태로 저장할 것이다.

dict를 이용해 저장하면 다음과 같이 저장할 수 있을 것이다.

# LEFT dict
{0: 1, 1: 1, 3: 2, 6: 3, 10: 3, 15: 1, 11: 2, 7: 3, 12: 2, 8: 3, 4: 2, 13: 1, 9: 3, 5: 3, 2: 1, 14: 1}
# RIGHT dict
{0: 1, -1: 1, -3: 2, -6: 3, -10: 3, -15: 1, -11: 2, -7: 3, -12: 2, -8: 3, -4: 2, -13: 1, -9: 3, -5: 3, -2: 1, -14: 1}

이때 숫자 0이 key값으로 들어갔는데, 0은 공집합이라고 생각하면 된다.

 

그러면 위 딕셔너리가 의미하는 것은 다음과 같게 된다.

딕셔너리의 key: 리스트에서 만들어낼 수 있는 부분수열의 합

딕셔너리의 value: 부분수열의 합을 만들어낼 수 있는 모든 경우의 수

 

이제 우리는 생각할 수 있다.

"LEFT dict의 key값을 순회할 때 (S - left key) 가 right key 안에 존재한다면, S를 만들 수 있는 경우의 수를 구할 수 있다"

이때 S를 만들 수 있는 경우의 수는 (left dict value) + (right dict value) 값이 될 것이다.

 

단, S가 0인 경우에는 left와 right 양쪽에 공집합이 하나씩 존재하므로 전체 count에서 1을 빼주어야 한다.

 

코드

LEFT_DICT = dict()
RIGHT_DICT = dict()


# seq: 부분수열, lst: 원본 리스트, sum_v: 누적합, idx: 현재 인덱스, option: left | right
def Backtracking(seq, lst, sum_v, idx, option):
    if option == 'LEFT':
        if sum_v in LEFT_DICT.keys():
            LEFT_DICT[sum_v] += 1
        else:
            LEFT_DICT[sum_v] = 1
    elif option == 'RIGHT':
        if sum_v in RIGHT_DICT.keys():
            RIGHT_DICT[sum_v] += 1
        else:
            RIGHT_DICT[sum_v] = 1

    for i in range(idx, len(lst)):
        next_lst = seq + [lst[i]]
        next_sum = sum_v + lst[i]
        Backtracking(next_lst, lst, next_sum, i+1, option)
        next_lst.pop()


N, S = map(int, input().split())
arr = list(map(int, input().split()))

start = 0
mid = len(arr)//2
end = len(arr)

Backtracking([], arr[start:mid], 0,  0, 'LEFT')
Backtracking([], arr[mid:end], 0, 0, 'RIGHT')

count = 0
for left in LEFT_DICT.keys():
    if (S - left) in RIGHT_DICT.keys():
        count += LEFT_DICT[left] * RIGHT_DICT[(S - left)]

if S == 0:
    count -= 1
print(count)

#백준 #0000번 #문제이름 #파이썬