[골드 3] 백준 2629 - 양팔저울 (파이썬)

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

풀이

• 추는 영쪽 저울에 올릴 수 있다.
• 이떄 만들 수 있는 무게 차를 추들을 사용해서 계산한다.
• D[i][j]는 앞에서 i개의 추를 사용해서 j - zero 만큼의 무게 차이를 만들 수 있는가?

N = int(input())  # 추의 개수
A = list(map(int, input().split()))  # 추의 무게 리스트
M = int(input())  # 구슬의 개수
B = list(map(int, input().split()))  # 구슬의 무게 리스트

zero = 40001로 초기화 해주었다. 그 이유는 인덱스에서 음수를 표현하기 위해서 기준점 0을 40001로 설정하였다.

점화식의 초기 값을 다음과 같이 설정한다.

D[0][zero] = True                         # 아무 것도 올리지 않은 경우
D[0][zero + A[0]] = True                 # 오른쪽에 A[0] 올린 경우
D[0][zero - A[0]] = True                 # 왼쪽에 A[0] 올린 경우

점화식은 아래와 같이 구성한다.

for i in range(1, N):
    for j in range(0, 80002):
        D[i][j] = D[i - 1][j]  # 현재 추를 안 쓴 경우

        if j - A[i] >= 0:
            D[i][j] |= D[i - 1][j - A[i]]  # 오른쪽에 A[i] 올린 경우
        if j + A[i] < 80002:
            D[i][j] |= D[i - 1][j + A[i]]  # 왼쪽에 A[i] 올린 경우

각 추는 저울의 왼쪽, 오른쪽, 사용 안 함의 세가지 선택지를 가진다.

위 연산은 기존에 만들 수 있던 무게차에서 +-A[i]를 더해 새 무게차도 만들수 있다는 의미이다.

for i in range(M):
    if D[N - 1][zero + B[i]]:
        print("Y")
    else:
        print("N")

구슬 무게 B[i]가 D[N - 1][zero + B[i]]에서 True이면 만들 수 있는 무게이므로 Y를 출력한다.

시간 복잡도 분석

• 공간 복잡도 : O(N * 80002)
• 시간 복잡도 : O(N * 80002)

코드

# 백준 2629 - 양팔저울
# 분류 : 다이나믹 프로그래밍

N = int(input())
A = list(map(int, input().split()))
M = int(input())
B = list(map(int, input().split()))

D = [[False] * 80002 for _ in range(N)]

# D[i][j] 0 ~ i -> j
# j < 0, 0 -> 40001

zero = 40001

D[0][zero] = True
D[0][zero + A[0]] = True
D[0][zero - A[0]] = True

for i in range(1, N) :
    for j in range(0, 80002) :
        D[i][j] = D[i - 1][j]

        if j - A[i] >= 0 :
            D[i][j] |= D[i - 1][j - A[i]]
        if j + A[i] < 80002 :
            D[i][j] |= D[i - 1][j + A[i]]

for i in range(M) :
    if D[N - 1][zero + B[i]]  :
        print("Y")
    else :
        print("N")