[백준] 5557번: 1학년- 파이썬

https://www.acmicpc.net/problem/5557

5557번: 1학년

문제

상근이가 1학년 때, 덧셈, 뺄셈을 매우 좋아했다. 상근이는 숫자가 줄 지어있는 것을 보기만 하면, 마지막 두 숫자 사이에 '='을 넣고, 나머지 숫자 사이에는 '+' 또는 '-'를 넣어 등식을 만들며 놀고 있다. 예를 들어, "8 3 2 4 8 7 2 4 0 8 8"에서 등식 "8+3-2-4+8-7-2-4-0+8=8"을 만들 수 있다.

상근이는 올바른 등식을 만들려고 한다. 상근이는 아직 학교에서 음수를 배우지 않았고, 20을 넘는 수는 모른다. 따라서, 왼쪽부터 계산할 때, 중간에 나오는 수가 모두 0 이상 20 이하이어야 한다. 예를 들어, "8+3+2-4-8-7+2+4+0+8=8"은 올바른 등식이지만, 8+3+2-4-8-7이 음수이기 때문에, 상근이가 만들 수 없는 등식이다.

숫자가 주어졌을 때, 상근이가 만들 수 있는 올바른 등식의 수를 구하는 프로그램을 작성하시오.

 

입력

첫째 줄에 숫자의 개수 N이 주어진다. (3 ≤ N ≤ 100) 둘째 줄에는 0 이상 9 이하의 정수 N개가 공백으로 구분해 주어진다.

 

출력

첫째 줄에 상근이가 만들 수 있는 올바른 등식의 개수를 출력한다. 이 값은 263-1 이하이다.

 

풀이

일단 "dfs로 풀면 되지 않을까?" 생각을 하고 풀었는데...

기다려도 예제2의 답을 출력을 못하는 것을 보고 다시 풀어야겠다고 생각했다.

 

스터디를 같이하는 팀원중 YG님이 하시던 말이 떠올랐다.

 

"일단 bfs, dfs로 접근해보고 뭔가 아닌가 싶으면 dp로 접근!"
- YG

 

예전에 풀었던 문제 중 기타리스트라는 dp문제가 있었는데, 그 때 곡의 볼륨을 리스트로 만들어서 접근하는 방법을 떠올리며 이 문제를 풀어봤다.

 

답안 코드

import sys
input = sys.stdin.readline

N = int(input())
NUM = list(map(int, input().split()))

dp = [[0]*(21) for _ in range(N)]
dp[0][NUM[0]]+=1

for idx in range(1, N-1):
    for result, count in enumerate(dp[idx-1]):
        if count:
            for new_result in (result+NUM[idx], result-NUM[idx]):
                if 0<=new_result<=20:
                    dp[idx][new_result] += count

print(dp[N-2][NUM[-1]])