문제

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

www.acmicpc.net/problem/1208

  • 입력

첫째 줄에 정수의 개수를 나타내는 N과 정수 S가 주어진다. (1 ≤ N ≤ 40, |S| ≤ 1,000,000) 둘째 줄에 N개의 정수가 빈 칸을 사이에 두고 주어진다. 주어지는 정수의 절댓값은 100,000을 넘지 않는다.

  • 출력

첫째 줄에 합이 S가 되는 부분수열의 개수를 출력한다.

과정

부분수열이래서 연결된 수열을 말하는 줄 알았는데 그건 아니었다. 그래서 combinations과 sum을 이용해서 풀었는데 바로 시간초과~.

import sys
from itertools import combinations

n, s = map(int, sys.stdin.readline().split())
num = list(map(int, sys.stdin.readline().split()))
res = 0

for x in num:
    if x == s:
        res += 1
for i in range(2, n):
    for comb in combinations(num, i):
        if sum(comb) == s:
            res += 1

print(res)

최종 코드

다른 방법이 도저히 생각이 안 나서 검색해봤는데 처음 보는 방법이었다.

수열을 left, right로 나눠서 나올 수 있는 수열의 합을 전부 구한 뒤, 사전에 저장한다. 그리고 사전의 key 값을 오름차순으로 정렬해서 left는 처음부터, right는 마지막부터 이분탐색한다.

이때 두 값의 합이 s라면 두 사전의 value 값을 곱해서 경우의 수를 더해준다. 만약 합을 리스트에 저장했다면 이렇게 바로 곱해서 경우의 수를 구할 수 없고 while 문을 통해 합이 같은 경우의 개수를 또 각각 구해서 곱해주어야 하는 번거로움이 생긴다. 그래서 사전을 사용한다.

그렇지 않고 합이 s보다 작다면 index l의 값을 1 증가시켜 합을 키우고, 합이 s보다 크다면 index r의 값을 1 감소시켜 합을 줄인다.

이 과정을 l < len(lkey) and r >= 0일 동안 반복하면 답을 찾을 수 있다. 이때 주의할 점은 s가 0일 때 수열의 크기가 0인 경우까지 count되므로 답에서 꼭 1을 빼주어야 한다.

import sys
from itertools import combinations

n, s = map(int, sys.stdin.readline().split())
num = list(map(int, sys.stdin.readline().split()))
res = 0

left = num[:n//2]
right = num[n//2:]
lsum = {0:1}
rsum = {0:1}
# 수열 왼쪽에서 나올 수 있는 모든 경우의 수 구하기
for i in range(1, len(left)+1):
    for comb in combinations(left, i):
        tmp = sum(comb)
        # 이미 사전에 있는 경우 +1
        if lsum.get(tmp):
            lsum[tmp] += 1
        # 없는 경우 value 값 1로 생성
        else:
            lsum[tmp] = 1
# 수열 오른쪽에서 나올 수 있는 모든 경우의 수 구하기
for i in range(1, len(right)+1):
    for comb in combinations(right, i):
        tmp = sum(comb)
        if rsum.get(tmp):
            rsum[tmp] += 1
        else:
            rsum[tmp] = 1

# key 값만 모아서 정렬
lkey = sorted(lsum.keys())
rkey = sorted(rsum.keys())
# lkey는 처음부터, rkey는 마지막부터 탐색
l, r = 0, len(rkey) - 1
while l < len(lkey) and r >= 0:
    tmp = lkey[l] + rkey[r]
    # 합이 s인 경우 두 경우의 수를 곱해서 최종 값에 더해준다
    if tmp == s:
        res += lsum[lkey[l]] * rsum[rkey[r]]
        l += 1
        r -= 1
    elif tmp < s:
        l += 1
    else:
        r -= 1

# s가 0인 경우 수열의 크기가 0인 경우까지 count되므로 1을 빼준다
if s == 0:
    res -= 1

print(res)

참고 사이트