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

참고 블로그:

https://ku-hug.tistory.com/190

[Python/파이썬] 백준 1208번: 부분수열의 합 2

 

깃허브:

https://github.com/MSIQOC/BOJ/blob/master/b1208_%EB%B6%80%EB%B6%84%EC%88%98%EC%97%B4%EC%9D%98%ED%95%A92.py

 

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

1208번: 부분수열의 합 2

 

여담을 하나 하자면 사실 나는 개발자로서의 길을 계속 걸어가다가 나랑은 안맞을 것 같다는 생각과 너무 늦게 시작한 것 같다는 생각에 교육대학원을 신청해두고 교사로서의 길도 생각하는 중이다. 그래서 이번에도 설명을 좀 자세하게 준비했는데 시작해보도록 하겠다...ㅎㅎ

 

먼저 전체적인 과정은 다음과 같다.

1. 수열을 반으로 쪼갠다.

2. 쪼갠 각각에 대해 부분수열의 합을 구해준다.

3. 이분탐색을 활용해서 합이 s가 되는 부분수열의 개수를 찾아준다.

 

먼저 주어진 수열을 arr에 저장했을 때 파이썬의 n//2를 활용해서 left = [:n//2], right = [n//2:]로 arr을 left와 right로 쪼개준다.

left의 부분수열, right의 부분수열을 구해준 다음에 각각의 부분수열의 합을 leftSum과 rightSum에 저장시키는 단계이다.

예를 들어서 현재 left가 [-7, -3]인데 부분수열을 구하면 [-7], [-3], [-7, -3]이 된다. 여기에서 각각의 합을 구하면

[-7] => -7

[-3] => -3

[-7, -3] => -7+(-3) = -10

이기 때문에 결국 leftSum에는 [-7, -3, -10]이 저장된다.

rightSum도 똑같은 방식으로 구해주면 된다.

 

그 다음은 사실상 leftSum에서 원소 하나를 선택하고 rightSum에서 원소 하나를 선택해서 두개의 합이 0이 되는 개수를 찾으면 된다.

여기에서 주의해야할 점은 예제에서는 n이 5여서 브루트포스로 leftSum의 원소와 rightSum의 원소를 하나씩 매치시키면서 합이 s가 되는 개수를 찾아도 문제가 없지만 n은 최대 40이다.

n이 최대 40이란 것은 n이 40일 때 leftSum의 원소가 20개, rightSum의 원소가 20개인데...

그렇게 되면 브루트포스로 했을 때 2^20 * 2^20 = 2^40으로 1조가 돼서 무조건 시간초과가 난다.

그러면 어떻게 하면 좋을까?

3번 과정은 3개의 스텝으로 다시 나눴다.

 

먼저 leftSum에 있는 모든 원소를 돌면서 l을 선택해준다.

-7은 이미 선택하고 넘겼다고 가정하고 l이 -3이라고 해보자.

여기에서 s-l은 3이 되는데, 이 때 이 3을 rightSum에서 이분탐색으로 찾아주고 존재한다면 answer에 1을 추가해주면 된다.

 

먼저 rightSum을 오름차순으로 정렬해주면 다음과같이 나온다.

여기에서 bisect_right, bisect_left는 이분탐색을 활용한 도구이다. 만약 위에서 r이 3이라면, bisect_left는 r이 시작하는 인덱스를 알려줘서 1을 리턴하고 bisect_right는 r이 끝나는 인덱스를 알려줘서 2를 리턴해준다.

그렇게 되면 위에서 bisect_right(rightSum, r) - bisect_left(rightSum, r)은 1이 된다.

 

여담으로 저기에서 3이 하나 더 있는 경우, 즉 rightSum이 [-2, 3, 3, 4, 5, 6, 8, 11, 13]인 경우, rightSum은 3이 끝나는 인덱스를 알려주기 때문에 3을 리턴하고 leftSum은 3이 시작하는 인덱스를 리턴해서 1을 리턴한다. 두개의 bisect의 차이를 구한 값은 2가 된다.

 

다시 구현했을 때 이 부분을 구현안해서 틀렸다ㅠㅠ

마지막으로 leftSum나 rightSum 자체도 부분집합의 합으로 이뤄져있기 때문에 그 안에서 s가 있는 경우도 계산을 해줘야한다.

count함수를 사용해도 되고 bisect함수를 사용해도 된다.

 

---

Python

"""
참고: https://ku-hug.tistory.com/190
"""
 
#함수 이름은 이해하기 쉽도록 작성했습니다.
 
import sys
from itertools import combinations
from bisect import bisect_left, bisect_right
input = sys.stdin.readline
 
def getSumOfArrayCombination(array, newArray):
    for i in range(1, len(array)+1):
        for a in combinations(array, i):
            newArray.append(sum(a))
 
 
n, s = map(int, input().split())
arr = list(map(int, input().split()))
left, right = arr[:n//2], arr[n//2:]
leftSum, rightSum = [], []
getSumOfArrayCombination(left, leftSum)
getSumOfArrayCombination(right, rightSum)
 
answer = 0
rightSum.sort()
 
for l in leftSum:
    r = s - l
    answer += bisect_right(rightSum, r) - bisect_left(rightSum, r)
 
answer += leftSum.count(s)
answer += rightSum.count(s)
 
print(answer)