[백준 / BAEKJOON] 1208번 부분수열의 합2 - Python

문제 원문

조건

시간 제한 : 1초
메모리 제한 : 256 MB

문제

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

입력

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

출력

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

예제

풀이

이 문제의 경우 주어지는 정수의 개수가 최대 40개 이므로 나올 수 있는 조합의 수를 계산하면 2^40이 됩니다.
(이항정리의 파생 공식으로 nC₀ + nC₁ + nC₂ + nC₃ + nC₄ + nC₅....  + nCn = 2^n)

2^40는 약 1조가 넘는 숫자로 모든 조합의 수를 계산하면 시간이 초과됩니다.

따라서 이 문제는 주어진 배열을 절반으로 나누어 각각 조합의 합을 구하고 구해진 두 배열의 합끼리 더하여 목표 값이 있는지 확인해야합니다.

이 때도 그냥 반복문으로 확인하면 시간이 오래걸리기 때문에 투 포인터를 활용한 이분탐색으로 확인합니다.

풀이과정은 다음과 같습니다.

1. 입력된 배열을 절반으로 나눈다 (왼쪽, 오른쪽)
2. 각각의 배열에서 나올 수 있는 조합의 합을 딕셔너리에 저장한다.
3. 각각의 딕셔너리에 목표 값이 있는지 확인한다. (오른쪽 인자을 포함하지 않거나 왼쪽 인자을 포함하지 않는 경우)
4. 각각의 딕셔너리의 value(합한 결과)를 정렬하여 리스트로 저장한 후 왼쪽은 index를 증가시키면서 오른쪽은 index를 감소시키면서 탐색한다.
5. 이때 각 배열이 포함되지 않는 경우는 3번 과정에서 확인 하였으므로 두 인덱스 중 하나라도 범위에서 벗어나면 반복을 종료한다.
6. 탐색 중 목표값과 일치하는 경우가 있다면 왼쪽의 해당 수의 갯수와 오른쪽의 해당 수의 갯수를 곱하여 결과에 더해준다.

성공 코드

import string
import sys

sys.setrecursionlimit(10**9)
input = sys.stdin.readline

arrSize, target = map(int, input().split())
fullArr = list(map(int, input().split()))
leftArr = fullArr[:arrSize//2]
rightArr = fullArr[arrSize//2:]
count = 0

#크기와 합한 결과만 넘긴다
def BT(sumDict, numArr, count = 0, lastIndex = -1, sumResult = 0):
    if count != 0:
        if sumResult in sumDict:
            sumDict[sumResult] += 1
        else:
            sumDict[sumResult] = 1

    for index in range(lastIndex + 1, len(numArr)):
        sumResult += numArr[index]
        BT(sumDict, numArr, count + 1, index, sumResult)
        sumResult -= numArr[index]

#왼쪽과 오른쪽 배열의 조합의 합을 계산한다
leftSum = dict()
rightSum = dict()
BT(leftSum, leftArr)
BT(rightSum, rightArr)

#왼쪽이나 오른쪽에서 해당하는 값이 있는지 찾아서 결과에 반영
if target in leftSum:
    count += leftSum[target]

if target in rightSum:
    count += rightSum[target]

#왼쪽과 오른쪽을 더해서 목표값이 되는 경우를 이분탐색으로 찾아 결과에 반영
leftResult = sorted(leftSum.keys())
rightResult = sorted(rightSum.keys())

leftIndex = 0
rightIndex = len(rightResult) - 1

while leftIndex < len(leftResult) and rightIndex >= 0:
    sumVal = leftResult[leftIndex] + rightResult[rightIndex]
    if sumVal == target:
        count += leftSum[leftResult[leftIndex]] * rightSum[rightResult[rightIndex]]

    if sumVal < target:
        leftIndex += 1
    else:
        rightIndex -= 1
    
print(count)

결과

Github - github.com/ledpear