[Java+Python]1463번, 쉬운 계단 수

문제

 

45656이란 수를 보자.

이 수는 인접한 모든 자리의 차이가 1이다. 이런 수를 계단 수라고 한다.

N이 주어질 때, 길이가 N인 계단 수가 총 몇 개 있는지 구해보자. 0으로 시작하는 수는 계단수가 아니다.

 

입력

 

첫째 줄에 N이 주어진다. N은 1보다 크거나 같고, 100보다 작거나 같은 자연수이다.

 

출력

 

첫째 줄에 정답을 1,000,000,000으로 나눈 나머지를 출력한다.

 

풀이

 

반복되는 동적계획법 문제이다.

 

계속 풀다 보니 동적계획법은 계획을 세우는 자체보다 경계선을 올바로 잡는 게 중요하다는 생각이 든다.

 

아니 어쩌면 경계선을 잡는 것 자체가 계획의 전부인지도.

 

우선 입력으로 주어지는 n을 받아 이차원 배열 memo를 선언했다.

 

그리고 0을 제외한 memo[1]의 요소를 1로 초기화했는데, 이는 문제의 조건에서 0으로 시작하는 수를

 

세지 말라고 지시했기 때문이다.

 

그다음은 이중반복문을 돌며 메모이제이션을 반복하는 것만 남았다.

 

mod(= 1_000_000_000)로 나누는 건 매 연산마다 반복해 줬으며

 

반복문이 끝나고 0에서 9까지 각 숫자로 끝나는 계단수의 개수를 mod로 나눈 값을 result에 전부 더해 출력해 주었다.

 

먼저 자바 코드다.

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;

public class Prob10844 {

	public static void main(String[] args) throws IOException {

		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

		int n = Integer.parseInt(br.readLine());
		int[][] memo = new int[n + 1][10];
		int mod = 1_000_000_000;
		int result = 0;

		for (int i = 1; i <= 9; i++) {
			memo[1][i] = 1;
		}

		for (int i = 2; i <= n; i++) {
			for (int j = 0; j <= 9; j++) {

				if (j == 0) {
					memo[i][0] = memo[i - 1][1] % mod;
				} else if (j == 9) {
					memo[i][9] = memo[i - 1][8] % mod;
				} else {
					memo[i][j] = (memo[i - 1][j - 1] + memo[i - 1][j + 1]) % mod;
				}
			}
		}

		for (int i = 0; i <= 9; i++) {
			result = (result + memo[n][i]) % mod;
		}

		System.out.println(result);
	}
}

다음은 파이썬 코드이다.

 

이번에는 파이썬이 메모리 두 배, 실행속도는 1/3을 기록했다.

import sys

n = int(sys.stdin.readline().rstrip())
memo = [[0 for _ in range(10)] for _ in range(n + 1)]
mod = 1_000_000_000
result = 0

for i in range(1, 10):
    memo[1][i] = 1

for i in range(2, n + 1):
    for j in range(10):
        if j == 0:
            memo[i][0] = memo[i - 1][1] % mod

        elif j == 9:
            memo[i][9] = memo[i - 1][8] % mod

        else:
            memo[i][j] = (memo[i - 1][j - 1] + memo[i - 1][j + 1]) % mod

for i in range(10):
    result = (result + memo[n][i]) % mod

print(result)