백준 5557번 : 1학년 in Python

5557번: 1학년

 

www.acmicpc.net

이 문제는 0 이상 9 이하의 정수 N개가 주어졌을 때 마지막 두 숫자 사이에 등호를 넣고, 나머지 숫자 사이에는 더하기 기호나 빼기 기호를 넣어 등식을 완성시키는 방법의 수를 구하는 문제이다. 단, 중간 연산 결과가 0보다 작거나 20보다 크면 안 된다.

만약 마지막 조건, 즉 중간 연산 결과가 0보다 작거나 20보다 크면 안 된다는 조건이 없으면 두 번째 정수부터 N-1번째 정수 각각에 대해 더하거나 빼는 모든 경우의 수를 고려해야 한다. 그러므로 결과를 반환하는데 걸리는 시간이 O(2^N)이나 걸리게 된다. 

하지만 해당 조건 덕분에 이 문제는 dynamic programming을 이용해 쉽게 해결할 수 있다. 우선 N*21 크기의 행렬 dp를 만든다. 여기서 dp의 (i, j) 성분은 i번째 수까지 고려했을 때 연산 결과가 j인 경우의 수를 의미한다. 그러므로 초기값으로 dp의 (0, (첫 번째 정수)) 성분을 1로 하고 나머지 성분을 0으로 한다.

그 다음 i=1부터 i=N-2까지 순회하면서 (i+1)번째 정수 num을 0~20 사이의 숫자 j에 더하거나 빼서 그 결과가 0~20 사이에 있으면 dp의 (i, j+num) (num을 더하는 경우) 또는 (i, j-num) (num을 빼는 경우) 성분에 dp의 (i, j) 성분을 더한다. 순회를 마치면 dp의 (N-2, (N번째 정수)) 성분을 통해 연산 결과가 N번째 정수와 같은 경우의 수를 얻을 수 있다.

이를 코드로 구현하면 다음과 같다.

N = int(input())
ints = list(map(int, input().split()))
dp = [[0]*21 for _ in range(N)]
# 초기값 설정
dp[0][ints[0]] = 1
for i in range(1, N-1):
    num = ints[i]
    # 0~20 사이의 모든 수를 고려한다.
    for j in range(21):
        # i+1번째 수에 해당하는 수를 더했을 때 
        if j+num <= 20:
            dp[i][j+num] += dp[i-1][j]
        # i+1번째 수에 해당하는 수를 뺐을 때
        if j-num >= 0:
            dp[i][j-num] += dp[i-1][j]
print(dp[N-2][ints[N-1]])