백준 2629번 : 양팔저울 in Python

2629번: 양팔저울

 

www.acmicpc.net

이 문제는 여러 가지 추들을 이용해 주어진 구슬의 무게를 확인할 수 있는지 결정하는 문제로, 단순히 모든 추에 대해 추를 저울 위에 올리지 않거나, 왼쪽 저울에 올리거나, 오른쪽 저울에 올리는 경우로 나눠 계산하는 경우 추의 개수가 n일 때 O(3^n)의 시간 복잡도를 가지므로, brute force로 계산하는 것은 매우 오래 걸린다. 그래서 이 문제는 다른 방법으로 풀어야 하는데, 바로 dynamic programming을 이용해 해결한다.

우선 처음에 생각한 방법은 최대 확인 가능한 구슬의 무게가 40000이므로 길이가 40001인 배열 dp를 만들어 추를 하나씩 추가하면서 확인 가능한 구슬의 무게에 해당하는 인덱스의 값을 1로 하는 방법이었다. 추를 하나씩 추가하면서 일단 추 자체의 무게에 해당하는 인덱스의 값을 1로 하고,  dp를 순회하면서 값이 1인 entry의 인덱스에 대해 추의 무게를 뺴고 더한 인덱스에 해당하는 값을 1로 만드는 방법을 생각했다. 그러나 이 방법은 문제점이 두 가지가 있었다. 먼저 추를 가벼운 것부터 골라 확인 가능한 구슬 무게를 따지면 무거운 추를 고르면서 추의 무게를 뺀 결과가 마이너스가 되는 경우를 고려하지 않는 것이었다. 그보다 더 심각한 문제는, dp를 순회하면서 추의 무게를 더한 인덱스에 대해 다시 추의 무게를 더한 인덱스에 해당하는 값을 1로 만들면서 엉뚱한 결과가 나온 것이다.    

그래서 앞의 문제를 해결하기 위해 무거운 추부터 따지기 시작해 추의 무게를 뺀 결과가 마이너스가 되지 않도록 했고(마이너스가 된 경우는 무시했다.), 뒤의 문제를 해결하기 위해 확인 가능한 구슬의 무게 집합을 도입했다.  그렇게 함으로써 무거운 추부터 순회하며 확인 가능한 구슬의 무게 집합에 대해 순회하며 확인 가능한 구슬의 무게를 점점 늘려나갔다. 여기까지 했을 때 예제 케이스도 맞아서 괜찮을 줄 알았는데, 실제로 제출을 하니 27% 정도에서 틀렸다고 나왔다. 

왜 틀렸나 생각을 해 본 결과 앞의 문제를 완전히 해결하지 못했다는 것을 깨달았다. 왜냐하면 추의 무게를 뺀 결과가 마이너스인 경우도 무시할 수 없기 때문이다. 예를 들어 7g, 4g, 2g, 2g, 2g의 추를 순서대로 추가한다고 할때, 네 번째 추까지 추가하는 경우 7g-4g-2g-2g = -1g이 나와 원래는 이를 무시했는데, 만약 마지막 2g의 추를 추가하면 7g-4g-2g-2g+2g = 1g이 나와 이는 실제로 가능한 경우의 수가 된다. 물론 1g은 7g-4g-2g=1g을 통해서도 구할 수 있지만, 추 무게가 어떤지에 따라 이렇게 구한 무게 역시 무시할 수 없다. 그래서 집합에 마이너스 값 역시 추가하도록 했고, 만약 이 마이너스 값이 어떤 추의 무게를 더해 양수가 되는 경우 배열 dp에 반영하도록 했다. 그랬더니 모든 테스트 케이스를 통과했다.

이를 구현한 코드는 다음과 같다.

N = int(input())
weights = list(map(int, input().split()))
dp = [0]*40001
# 확인 가능한 구슬의 무게 집합
s = set()
# 무거운 추부터 순회한다.
for i in range(N-1, -1, -1):
    weight = weights[i]
    dp[weight] = 1
    update_s = []
    # 각 확인 가능한 무게에 추의 무게를 빼고 더해 집합에 추가하고, 특히 그 결과가 양수면 dp에 반영한다.
    # 집합을 순회하는 동안 집합의 크기가 변하면 안 되므로 
    # 새 배열에 추가할 무게를 모두 추가한 뒤 집합에 반영했다.
    for w in s:
        update_s.append(w-weight)
        update_s.append(w+weight)
        if w - weight > 0:
            dp[w - weight] = 1
        if 0 < w + weight <= 40000:
            dp[w+weight] = 1
    s.update(update_s)
    # 기본 추의 무게는 마지막에 추가해 추 무게의 2배에 해당하는 무게가 불필요하게 반영되는 것을 막는다.
    s.add(weight)
M = int(input())
marbles = list(map(int, input().split()))
for i in range(M):
    if i == M-1:
        if dp[marbles[i]] == 1:
            print("Y")
        else:
            print("N")
    else:
        if dp[marbles[i]] == 1:
            print("Y", end=" ")
        else:
            print("N", end=" ")