[Python] 백준 1208번: 부분수열의 합2

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

1208번: 부분수열의 합 2

 

import sys
import itertools
input = sys.stdin.readline


def setValidSetCount():
    global validSetCount

    leftSums = []
    for i in range(N//2 + 1):
        for s in itertools.combinations(nums[:N//2], i):
            leftSums.append(sum(s))
    leftLength = len(leftSums)
    leftSums.sort()

    rightSums = []
    for i in range(N - N//2 + 1):
        for s in itertools.combinations(nums[N//2:], i):
            rightSums.append(sum(s))
    rightLength = len(rightSums)
    rightSums.sort(reverse=True)

    l, r = 0, 0
    while l < leftLength and r < rightLength:
        curLeft, curRight = leftSums[l], rightSums[r]
        curSum = leftSums[l] + rightSums[r]

        if curSum == S:
            ls, rs = l, r
            while ls < leftLength and leftSums[ls] == curLeft:
                ls += 1
            while rs < rightLength and rightSums[rs] == curRight:
                rs += 1

            validSetCount += (ls-l)*(rs-r)
            l, r = ls, rs
        elif curSum < S:
            l += 1
        elif curSum > S:
            r += 1


if __name__ == '__main__':
    N, S = map(int, input().split())
    nums = list(map(int, input().split()))

    validSetCount = 0
    setValidSetCount()

    print(validSetCount-1 if S == 0 else validSetCount)

 

문제는 간단하나, 시간복잡도와가 문제이다. N이 40일 때, itertools.combinations()을 돌리게 되면 절대로 풀 수 없다. 따라서 반으로 나누어서 고려한다. 반으로 나뉜 리스트 두 개에 대해서 combinations를 이용하여, 모든 경우의 합의 정보를 담는 리스트인 leftSums와 rightSums를 만들어준다. 예제 입력 1과 같은 경우에는 다음과 같이 나뉜다.

[0, -7, -3, -10], [0, -2, 5, 8, 3, 6, 13, 11]

여기서 앞에 0이 있는데, 이것은 나중에 설명하겠다. 그 다음에는 두 리스트를 각각 오름차순, 내림차순으로 정렬시킨다.

[-10, -7, -3, 0], [13, 11, 8, 6, 5, 3, 0, -2]

이제, 각각의 리스트의 인덱스를 가리키는 변수 2개(l, r)를 통해 앞에서부터 탐색해갈 것이다. l과 r이 가리키는 각 리스트의 값을 합한 값이 curSum인데, 이것을 입력받은 S와 비교해준다. 만약 같을 시에는, 중복될 수도 있음을 감안하여, 값이 바뀌기 전까지 각각의 리스트를 탐색한다. 이에 따른 결과인 ls, rs를 통해 validSetCount와 l, r에 적절한 값을 넣어준다. 그리고, 2개의 elif문은, l을 증가시키면 curSum이 감소하고, r을 증가시키면 curSum이 감소한다는 것을 감안하면 쉽게 생각해낼 수 있다.(투포인터)

마지막으로, 이렇게 validSetCount 값이 변동되면 이것을 출력해주면 되는데, 조건이 하나 있다. 만약 S가 0일 때는 1을 빼주어야 한다는 것이다. 다시 돌아가서, leftSum과 rightSum에 0이 있는 이유는, S와 조건이 부합하는 sum값을 만들어내는 set이 한쪽의 리스트에만 포함되어 있을 수도 있기 때문이다. 예를 하나 들자면, 만약 0을 포함시키지 않아서 leftSum이 [-10, 10]이고, rightSum이 [99, 100]이라고 하면 결과값이 1이 나오지 않는다. 아무튼, 이런 이유로 leftSums, rightSums에 combinations로 값을 추가해나갈 때 for문의 범위를 0부터 시작하도록 해준 것이다. 그렇다면, S가 0일 때는 왜 1을 빼주어야할까? 그것은 바로 예시 입력 1에서 바로 알 수 있는데, 0이라는 것이 위의 이유 때문에 만들어졌으나, S가 0일 때는 그렇게 만들어진 0과 0이 결과값에 영향을 1번 주기 때문이다.