[백준/2629] 양팔 저울 (Python)

 

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

2629번: 양팔저울

 

Difficulty : Gold 3

 

Status : Solved

 

Time : 00:11:50

 

---

 

문제 설명

 

양팔 저울과 몇 개의 추가 주어졌을 때, 이를 이용하여 입력으로 주어진 구슬의 무게를 확인할 수 있는지를 결정하려고 한다.

무게가 각각 1g과 4g인 두 개의 추가 있을 경우, 주어진 구슬과 1g 추 하나를 양팔 저울의 양쪽에 각각 올려놓아 수평을 이루면 구슬의 무게는 1g이다. 또 다른 구슬이 4g인지를 확인하려면 1g 추 대신 4g 추를 올려놓으면 된다.

구슬이 3g인 경우 아래 <그림 1>과 같이 구슬과 추를 올려놓으면 양팔 저울이 수평을 이루게 된다. 따라서 각각 1g과 4g인 추가 하나씩 있을 경우 주어진 구슬이 3g인지도 확인해 볼 수 있다.

<그림 2>와 같은 방법을 사용하면 구슬이 5g인지도 확인할 수 있다. 구슬이 2g이면 주어진 추를 가지고는 확인할 수 없다.

추들의 무게와 확인할 구슬들의 무게가 입력되었을 때, 주어진 추만을 사용하여 구슬의 무게를 확인 할 수 있는지를 결정하는 프로그램을 작성하시오.

 

입력 및 출력

 

입력

첫째 줄에는 추의 개수가 자연수로 주어진다. 추의 개수는 30 이하이다. 둘째 줄에는 추의 무게들이 자연수로 가벼운 것부터 차례로 주어진다. 같은 무게의 추가 여러 개 있을 수도 있다. 추의 무게는 500g이하이며, 입력되는 무게들 사이에는 빈칸이 하나씩 있다. 세 번째 줄에는 무게를 확인하고자 하는 구슬들의 개수가 주어진다. 확인할 구슬의 개수는 7이하이다. 네 번째 줄에는 확인하고자 하는 구슬들의 무게가 자연수로 주어지며, 입력되는 무게들 사이에는 빈 칸이 하나씩 있다. 확인하고자 하는 구슬의 무게는 40,000보다 작거나 같은 자연수이다.

 

출력

주어진 각 구슬의 무게에 대하여 확인이 가능하면 Y, 아니면 N 을 차례로 출력한다. 출력은 한 개의 줄로 이루어지며, 각 구슬에 대한 답 사이에는 빈칸을 하나씩 둔다.

 

입력 예시

2
1 4
2
3 2

 

출력 예시

Y N

 

 

---

 

풀이

DP 문제, 그 중에서도 배낭 문제의 확장판이라 볼 수 있는 문제이다.

 

현재 총 weight만큼 양팔저울에 올려져 있고 어떤 추 next_weight를 양팔저울에 넣을 때, 세 가지 가능성이 존재한다.

• weight 쪽에 추를 올리기 : 이 때 총 무게는 weight + next_weight가 된다.
• 아무 것도 하지 않기 : 이 때 총 무게는 weight로 변동이 없다.
• weight 반대쪽에 추를 올리기 : 이 때 총 무게는 weight - next_weight가 된다.

 

즉 세 가지 가능성이 발생한다.  또한 weight - next_weight 케이스의 경우 음수가 될 수 있다. 이 때는 양팔저울의 추를 서로 교환(swap)한다는 의미로, 절댓값을 취해 줄 수 있다. 따라서 i번째 추로 j만큼의 무게를 측정 가능한 경우(dp[i][j] = True), i+1번째 추를 이용하여 다음과 같은 점화식으로 표현 가능해진다.

 

DP[i+1][j + next_weight] = DP[i+1][j] = DP[i+1][abs(j+next_weight)] = True,
if Dp[i][j] == True

 

 

총 무게는 40000g 까지 가능하며, 추가 30개까지 존재할 수 있으므로 30 * 3 * 40000안에 dp를 작성 가능하다.

 

• init : 문제풀이에 필요한 정보를 입력받아 반환하는 함수이다.
• make_weight_dp : 위에서 표현한 점화식에 따라 2차원 리스트 DP를 생성하여 반환하는 함수이다.
• is_enable_measure_weight : 측정하고자 하는 무게 리스트 target_weight의 각 원소에 대해, DP의 그 무게에 해당하는 열을 검사하여 가능한지의 여부를 검사하는 함수이다. 검사 결과 리스트인 result 를 반환한다.
• solve : 메인함수. init 함수로 받은 정보를 make_weight_dp 함수와 is_enable_measure 함수로 전달하고, 그 결과 리스트 result를 출력한다.

 

풀이 코드

def init() :
  N = int(input())
  weight_list = list(map(int, input().split()))
  M = int(input())
  target_weight = list(map(int, input().split()))
  return N, M, weight_list, target_weight


def make_weight_dp(N, weight_list) :
  dp = [[False]*40001 for _ in range(N)]

  dp[0][0] = dp[0][weight_list[0]] = True

  for i in range(N-1) :
    next_weight = weight_list[i+1]
    for j in range(40001) :
      if dp[i][j] :
        for k in [-1, 0, 1] :
          dp[i+1][abs(j + k*next_weight)] = True

  return dp

def is_enable_measure_weight(N, M, target_weight, dp) :
  result = list()
  for i in range(M) :
    flg = False
    for j in range(N) :
      if dp[j][target_weight[i]] :
        result.append('Y')
        flg = True
        break
    if not flg : 
      result.append('N')

  return result

def solve() :
  N, M, weight_list, target_weight = init()
  dp = make_weight_dp(N, weight_list)
  result = is_enable_measure_weight(N, M, target_weight, dp)
  print(*result)

solve()

풀이 완료~~