(파이썬)백준 2629번 : 양팔저울

백준 2629번

임의의 무게를 주어진 추와 양팔 저울을 이용하여 측정할 수 있는 지 확인하는 문제입니다.
예를 들어, 1g과 4g의 추가 있을 때 잴 수 있는 무게는 1g, 3g, 4g, 5g입니다.

이번엔 M개의 무게추가 왼쪽 저울에 올라가 있고 무게는 100g이라고 할게요.
이 때 측정할 수 있는 무게는 100g입니다. 당연하죠?
저한테 5g짜리 추가 하나 더 있다고 해보죠. 그럼 어떤 구슬을 측정할 수 있죠?

첫번째 : 왼쪽에 추를 추가해서 105g을 측정할 수 있습니다.
            -> 100 + 5 = marble
두번째 : 추를 오른쪽에 95g짜리 구슬과 함께 올리면 수평을 이루겠죠?
            -> 100 = marble + 5
세번째 : 추를 아예 사용하지 않을 수도 있습니다.
            -> 100 = marble

결국 하나의 추를 더 사용할 때, 발생하는 경우의 수는 3가지인 것입니다.
주의할 점은 추를 사용한다는 것이, 반드시 저울에 올려야하는 것은 아니다 라는것이죠.

그렇다면 추의 개수를 하나씩 추가해가며 3가지 경우에 대해 확인해가며
마지막 추까지 진행하면 될 것 같습니다. 이젠 DP문제라는 것이 보이실겁니다.
바로 구현으로 넘어가볼게요.

1. 구현

늘 그렇듯 입력부터 짜보면
추의 개수 : n / 추 목록 : w_list / 구슬의 개수 : m / 구슬 목록 : m_list , 코드는

input = sys.stdin.readline

#추의 개수
n = int(input())
#추 오름차순 입력
w_list = list(map(int,input().split()))
#측정할 구슬의 개수
m = int(input())
m_list = list(map(int,input().split()))

DP문제이기때문에 DP배열을 사용할건데, 저는 2차원 배열을 사용하려 합니다.
DP[ 총 몇 개의 추를 사용했는 지 ][ 무게 ] 이런 형태로 말이죠.
무게가 측정가능한지 아닌지를 나타내야 하니까 초기값은 False로 하고 True로 갱신하겠습니다.

DP = [[False for _ in range(500 * n + 1)] for _ in range(n+1)]

위에서 말했듯, 추를 사용한다고해서 반드시 저울에 올릴 필요는 없습니다.
따라서 연산이 끝나고 나면 DP[ n ]에서 모든 측정 가능한 무게가 True로 갱신될겁니다.

이제는 함수를 구현해보겠습니다.
저는 재귀에 별로 자신이 없어서 반복문을 주로 이용하는 데
이 문제는 재귀로 푸는 게 좋아보입니다.

그렇다면 재귀 함수의 인자로는 어떤게 필요할까요?
DP배열을 차례로 갱신해가려면,
사용한 추의 개수가 몇 개인지와
현재 측정가능한 무게가 얼마인지를 인자로 받아야합니다.

def function(tmp,weight):

인자를 통해서
DP[tmp][weight]에 접근, 값을 True로 갱신하면,

DP[tmp][weight] = True

'tmp개의 추를 사용하여 weight라는 무게를 측정할 수 있다'
라는 뜻이 됩니다.

재귀적 호출은 어떻게 구현해야 할까요?
일단 위에서 말한 3가지 경우에 대해 적어보겠습니다.

1. 추를 왼쪽에 추가한다.
2. 추를 오른쪽에 추가한다.
3. 추를 사용하지 않는다.

왼쪽, 오른쪽보단 구슬이 없는 쪽, 있는 쪽이 적절한 표현이지만
편의상 왼쪽, 오른쪽으로 하겠습니다.

일단 3번은 쉽죠. 사용한 구슬의 개수만 1 추가하고
무게는 그대로입니다.

    #추 안쓰기
    function(tmp + 1, weight)

1번은 추를 왼쪽에 올리는 경우니까 추의 무게를 더하면 됩니다.

    #추 더하기
    function(tmp + 1, weight + w_list[tmp - 1])

2번은 빼는거죠? 음수가 되면 안되니까 절대값 처리를 해줘야합니다.

    #추 빼기
    function(tmp + 1, abs(weight - w_list[tmp - 1]))
    #빼는 경우는 절댓값 처리

인덱스만 제대로 잡으면 크게 어렵지 않을 것 같습니다.

구현도 거의 끝났습니다.
무한루프에 빠지지않게 제어문만 추가하면 됩니다.

1. 사용한 추의 개수가 n을 넘었을 때.
2. 이미 갱신한 값일 때.

코드로는

	#주어진 추의 개수를 넘어갔다면 #인덱스오류 방지
    if tmp > n :
        return
    #이미 계산했다면 #반복호출 방지
    if DP[tmp][weight]:
        return

이제 재귀함수를 초기값 0,0으로 실행시키고
문제의 출력 조건에 맞춰 결과를 출력하면 됩니다.
출력문은 생략하겠습니다.

전체코드

import sys
input = sys.stdin.readline

n = int(input())
w_list = list(map(int,input().split()))
m = int(input())
m_list = list(map(int,input().split()))

#배열 초기화, 초기값은 모두 False(불가능)
DP = [[False for _ in range(500 * n + 1)] for _ in range(n+1)]

# tmp = 몇번째 추에 대해서 계산하는지
# weight = 추를 넣냐 빼냐 안쓰냐에 따라 바뀌는 무게
def function(tmp,weight):
    #주어진 추의 개수를 넘어갔다면 #인덱스오류 방지
    if tmp > n :
        return
    #이미 계산했다면 #반복호출 방지
    if DP[tmp][weight]:
        return
    #가능으로 변경
    DP[tmp][weight] = True

    #다음 추로 넘어가며 3가지 경우에 대해 재귀호출

    #추 더하기
    function(tmp + 1, weight + w_list[tmp - 1])
    #추 빼기
    function(tmp + 1, abs(weight - w_list[tmp - 1])) #빼는 경우는 절댓값 처리
    #추 안쓰기
    function(tmp + 1, weight)

#main
function(0,0)
for m_weight in m_list:
    if m_weight > 500 * n:
        print("N",end=' ')
    else :
        if DP[n][m_weight]:
            print("Y",end=' ')
        else :
            print("N",end=' ')

질문과 피드백은 언제나 환영입니다.