[Python] 백준 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개가 공백으로 구분해 주어진다.

출력

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

정답코드

import sys

N = int(sys.stdin.readline())
numbers = list(map(int, sys.stdin.readline().split()))

dp = [[0] * 21 for _ in range(101)]
dp[0][numbers[0]] = 1

for i in range(1, N-1):
    for num in range(0, 21):
        if dp[i-1][num]:
            new_num = num + numbers[i]
            if 0 <= new_num <= 20:
                dp[i][new_num] += dp[i-1][num]
            new_num = num - numbers[i]
            if 0 <= new_num <= 20:
                dp[i][new_num] += dp[i-1][num]

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

첫 풀이

N = int(sys.stdin.readline())
numbers = list(map(int, sys.stdin.readline().split()))

dp = [[] for _ in range(N)]
dp[0] = [numbers[0]]
for i in range(1, N-1):
    for num in dp[i-1]:
        new_num = num + numbers[i]
        if 0 <= new_num <= 20:
            dp[i].append(new_num)
        new_num = num - numbers[i]
        if 0 <= new_num <= 20:
            dp[i].append(new_num)

answer = 0
for num in dp[N-2]:
    if num == numbers[-1]:
        answer += 1
print(answer)

-> 계산시간이 오래걸려서 아래 아이디어로 교체

dp를 2차원으로 구성, dp[idx][num] = {idx번째 계산결과에 해당 숫자가 몇개 있는지}
3 <= idx <= 100, 0 <= num <= 20

같은 숫자를 중복해서 dp에 추가하지 않고 해당 숫자가 몇 개있는지 체크