1 minute read

1208번 : 부분수열의 합 2

문제 링크

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

문제 설명

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

입력

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

출력

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

문제풀이

N의 개수가 최대 40개가 되므로 40개 항목의 부분배열을 구하기에는 O(2^40)으로 시간초과가 된다.

따라서 절반으로 나누어 20개의 항목의 부분배열을 2번 구하면 O(2^21)로 답을 구할 수가 있다.

  1. 일단 절반으로 나눈 배열 내에서 합이 S인 배열의 개수를 저장
  2. 절반으로 나눈 배열을 각각 A, B라고 했을 때 A[i] + B[j]가 S인 배열의 개수를 더한다.
from itertools import combinations

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

# 배열을 절반으로 나누기
a, b = arr[:N//2], arr[N//2:]
aSum, bSum = [], []

# 절반으로 나눈 배열 내 부분배열의 합들을 저장하기
for i in range(1, len(a)+1) :
    for comb in combinations(a, i) :
        aSum.append(sum(comb))
for i in range(1, len(b)+1) :
    for comb in combinations(b, i) :
        bSum.append(sum(comb))

# 절반으로 나눈 배열 내에서 합이 S인 배열의 개수를 저장 
answer = aSum.count(S) + bSum.count(S) 

# 각 배열을 오름차순 정렬, 이분탐색으로 하기 위해서
aSum.sort()
bSum.sort()

i, j = 0, len(bSum)-1

while i < len(aSum) and j >= 0 :
    tempSum = aSum[i] + bSum[j]

    # 만약 합이 S와 같다면 각 합배열 내에서도 같은 항목이 있는지 확인 후 그 개수를 곱해서 저장
    if tempSum == S :
        ti, tj = i, j

        while ti < len(aSum) and aSum[ti] == aSum[i] :
            ti += 1
        while tj >= 0 and bSum[tj] == bSum[j] :
            tj -= 1    

        answer += (ti-i) * (j-tj)
        i, j = ti, tj

    # 만약 합이 S보다 크다면 뒤 배열의 인덱스를 1 감소
    elif tempSum > S : 
        j -= 1

    # 만약 합이 S보다 작다면 앞 배열의 인덱스를 1 증가
    else : 
        i += 1

print(answer)

Comments

You May Also Enjoy