[Baekjoon]백준 2629 양팔저울(골드 3) - Python, C++

문제설명 1

문제설명 2

해당 문제는 해당 추의 개수를 가지고 원하는 무게를 만들수 있는지에 대해 물어보는 문제이다

예제 2번을 보면 구슬 4개를 가지고 1, 4, 10을 만들어야한다

1의 경우 3과 2, 4의 경우 3, 3과 2로 만들수 있고 10의 경우 2, 3, 3, 3의 추로는 만들수 없기 때문에 N이 출력된다

 

일단 dp라는 배열에 어떠한 값을 넣을지 고민해보았다 사실상 이것이 동적 프로그래밍의 핵심이다

dp[사용한 추의 개수][무게 차이]를 기준으로 1이라는 값을 넣어주고 나머지는 0으로 초기화한다

그렇게 하게 되면 함수를 재귀적으로 호출할 때 1인 경우 바로 return을 함으로써 불필요한 계산을 하지 않는다

또한 구슬이 기존에 입력된 수보다 더 큰 경우에도 바로 return을 하여준다

def sol(num, wg):
    if num > weight_num: # 구슬의 숫자가 더 큰 경우
        return
    if dp[num][wg] == 1: # 이미 연산함
        return
    dp[num][wg] = 1

이제부터 재귀적으로 호출하는 것의 핵심인데, 추를 추가하는 경우, 그대로 진행하는 경우

반대편에 추를 추가하는 경우 총 3가지이다(반대편에 추를 추가하는 경우는 무게를 빼는 과정과 같다)

무게를 뺄 경우 절댓값을 구해야하므로 이에 주의하여 인자로 전달하면 된다

sol(num + 1, wg + weight[num-1])       # 추 추가하기
sol(num + 1, wg)                       # 그대로 진행
sol(num + 1, abs(wg - weight[num-1]))  # 반대편에 추 추가(무게빼기)

해당 함수의 동작은 다음과 같고, 나머지 기타 조건(입력받기, 출력 등)은 밑에 전체코드에 있다

import sys
input = sys.stdin.readline

def sol(num, wg):
    if num > weight_num: # 구슬의 숫자가 더 큰 경우
        return
    if dp[num][wg] == 1: # 이미 연산함
        return
    dp[num][wg] = 1
    
    sol(num + 1, wg + weight[num-1])       # 추 추가하기
    sol(num + 1, wg)                       # 그대로 진행
    sol(num + 1, abs(wg - weight[num-1]))  # 반대편에 추 추가(무게빼기)

weight_num = int(input())
weight = list(map(int, input().split()))
bead_num = int(input())
bead = list(map(int, input().split()))
# dp[사용한 추의 개수][무게 차이]
dp = [[0 for _ in range(500 * weight_num + 1)] for _ in range(weight_num + 1)]

sol(0, 0)

for k in bead:
    if k > 500 * weight_num: # 가능한 구슬 무게 범위 초과
        print("N", end=' ')
    elif dp[weight_num][k]:  # 구슬 무게 측정 가능
        print("Y", end=' ')
    else:                    # 구슬 무게 측정 불가능
        print("N", end=' ')

문제조건에서 한 줄에 결과를 출력하라고 했기때문에 print이후에 end=' '를 해주어야한다

해당 코드를 C++로도 작성해보았는데 문제해결을 위한 함수에서 return할때 값을 전달하지 않으므로

void로 선언하여야하며 이외에는 파이썬과 동일하다

#include <iostream>
#include <cmath>
using namespace std;
int weight_num, bead_num;
int weight[31] = {0, };
int bead[8] = {0, };
// dp[사용한 추의 개수][무게 차이]
int dp[31][15001] = {0, };

void sol(int num, int wg) {
    if(num > weight_num) // 구슬의 숫자가 더 큰 경우
        return;
    if(dp[num][wg] == 1) // 이미 연산함
        return;
    dp[num][wg] = 1;

    sol(num + 1, wg + weight[num]);         // 추 추가하기
    sol(num + 1, wg);                       // 그대로 진행
    sol(num + 1, abs(wg - weight[num]));    // 반대편에 추 추가(무게빼기)
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);
    
    cin >> weight_num;
    for(int i = 0; i < weight_num; i++)
        cin >> weight[i];

    cin >> bead_num;
    for(int i = 0; i < bead_num; i++)
        cin >> bead[i];

    sol(0, 0);

    for(int i = 0; i < bead_num; i++) {
        if(bead[i] > 500 * weight_num)
            cout << "N" << " ";
        else if(dp[weight_num][bead[i]] == 1)
            cout << "Y" << " ";
        else
            cout << "N" << " ";
    }

    return 0;
}