[백준] 2629. 양팔 저울 (python)

문제 출처 : 2629번: 양팔저울
코드 참조 : https://source-sc.tistory.com/

문제 풀이

해당 문제는 동적 계획법 (Dynamic Programming)으로 푸는 문제이다.

저울의 사용자는 추를 가지고 세 가지 작업 중 하나를 수행할 수 있다.

1. 추를 왼쪽에 올린다.
2. 추를 오른쪽에 올린다.
3. 추를 올리지 않는다.

세 가지 작업을 모두 수행하는 DFS를 실행하여 문제를 해결한다.

코드

import sys
input = sys.stdin.readline

# 입력값 받기

num_pends = int(input()) # 추의 개수 (<= 30)
pends = list(map(int, input().split())) # 추 (오름차순, <= 500)
num_marbles = int(input()) # 구슬의 개수 (<= 7)
marbles = list(map(int, input().split())) # 구슬의 무게 (<= 40000)

# dp table 만들기

dp = [[0]*(30*500+1) for _ in range(num_pends+1)] 
# dp[i][j]  = i 번째 까지의 추를 놓았을 때, j 무게를 만들 수 있는지

result = set()

# 관련 함수 정의

def get_result(pends, n, now, left, right):
    """
    pends : 추 목록
    n : 전체 추의 개수
    now : 이제 놓을 추의 index
    left : 왼쪽 팔의 무게
    right : 오른쪽 팔의 무게
    """

    new = abs(left - right)
    # 처음 함수가 호출되었을 때의 무게를 잰다.

    if new not in result: # 처음 재는 무게라면
        result.add(new) # result에 추가

    if now == n : # 재귀 탈출 조건 (모든 추를 다 올려놓았을 때)
        return 0

    if dp[now][new] == 0: # now까지의 추로 아직 무게를 잰 적이 없다면
        get_result(pends, n, now+1, left + pends[now], right) # 왼쪽에 놓기
        get_result(pends, n, now+1, left, right + pends[now]) # 오른쪽에 놓기
        get_result(pends, n, now+1, left, right) # 놓지 않기
        dp[now][new] = 1 # 무게를 재었다고 표시

get_result(pends, num_pends, 0, 0, 0)
answer = []

for marble in marbles:
    if marble in result:
        answer.append("Y")
    else:
        answer.append("N")

print(*answer)

결과

배운 점

처음에 DP문제라는 것을 알았을 때, 추로 인해 계산된 무게를 저장해야한다는 강박관념에 빠져 문제를 풀지 못했다.

그러나 다른 사람의 코드를 보면서 "계산된 무게"를 저장하기 보다는, "무게가 계산되었는지 여부"를 저장하는 방법도 있다는 것을 배웠다.