백준 1562 - 계단 수

noj.am/1562

 

백준 1562번 계단 수 문제이다.

 

문제)

'계단 수'는 인접한 숫자간의 차가 모두 1인 수를 가리킨다.

자연수 N이 주어졌을 때, N자리 계단 수 중 0~9 모든 숫자가 등장하는 계단 수의 개수를 출력하세요.

 

풀이)

처음엔 N 이하인 모든 계단수의 개수를 계산하는 줄 알고 백트래킹을 생각했다. 그런데 문제를 다시 읽어보니 0~9 모든 숫자가 들어가야한다는 조건이 있어서 함수 인자로 used라는 set을 추가했다. 그럼에도 TLE(시간초과)가 나서 백트래킹으로 풀기 적합하지 않은 문제란 생각이 들었다. 백트래킹으로 계산하면 0~9가 모두 들어가지 못하는 케이스를 불필요하게 많이 체크하기 때문에 비효율이 발생한다.

 

문제에서 집중할 부분은 크게 3가지이다.

"지금까지 결정한 자리(digit)의 개수, 가장 마지막으로 고른 digit에 들어간 숫자, 지금까지의 (부분적인) 계단 수를 만드는 데 사용한 숫자의 종류"

다른 말로, 만든 수가 몇인지는 중요하지 않다는 뜻이다. 이는 문제의 제약 조건이 "N자리 수"이기 때문이다.

가장 마지막에 고른 digit에 들어간 숫자는 다음 digit을 고르기 위해 필요한 정보이다. 이를 알아야 계단 수를 만들 수 있기 때문이다.

지금까지의 계단 수를 만드는 데 사용한 숫자의 종류는 백트래킹 구현에서는 set 자료구조를 이용하였다. DP로 넘어오면서는 인덱스 형태로 표현할 수 있는 값이어야 해서 비트마스킹을 이용했다.

비트마스킹은 1<<i에 해당하는 비트가 1인지 0인지로 i에 관한 어떤 정보를 True/False로 나타내는 것을 의미한다.

 

내가 사용한 DP Table은 아래와 같다.

dp[i][j][k] : i자리 계단 수 중 마지막 수가 j로 끝나며, 사용한 숫자의 종류가 k로 비트마스킹 되는 계단 수의 개수

i가 가질 수 있는 값의 범위: 0 ~ N (N을 포함하기 위해 list는 N + 1 크기로 만들었다. 0-indexing으로 하는 경우 N 크기로 만듦.)

j가 가질 수 있는 값의 범위: 0 ~ 9 (따라서 크기 10)

k가 가질 수 있는 값의 범위: 0 ~ (1 << 10) - 1 (0은 아무것도 사용하지 않은 경우. 1 << 10 - 1은 모든 숫자를 사용한 경우)

 

이번 문제를 풀 때는 비트 연산 중 OR(|) 연산만 알면 된다.

| : 좌항과 우항의 비트를 비교하면서 한쪽이라도 켜진(1) 비트는 1로 계산하는 연산.

우리는 이것을 DP Table을 채울 때 새로운 digit의 사용을 k에 반영해주기 위해 활용한다.

 

전체적인 흐름은 다음과 같다.

1. 1자리 계단 수의 가짓수를 정리해준다(1~9까지 개수를 1로 설정한다)

2. 2~N자리 계단 수의 가짓수를 계산한다. (for i loop)

3. i - 1 자리 계단 수 중 마지막 digit이 j인 계단 수 중 (for j loop)

4. 사용한 digit의 종류가 비트마스킹 결과 k로 나타나는 수에 대해서 (for k loop)

5 - 1. j가 0 초과라면, dp[i][j - 1][k | (1 << (j - 1))] += dp[i - 1][j][k]

5 - 2. j가 9 미만이라면, dp[i][j + 1][k | (1 << (j + 1))] += dp[i - 1][j][k]

6. N자리 계단 수 중 index k가 MAX - 1인 경우의 가짓수를 모두 더한다.

7. 결과값을 1e+9로 나눈 나머지를 구한다.

8. %= 연산 결과 파이썬에서 Float으로 나온 값을 round 처리한다(int로 할 경우 부동소수점 오차에 의해 값이 살짝 작아지는 경우에 대해 오답이 생긴다. int는 버림하기 때문이다).

9. 8에서 나온 결과값을 출력한다.

 

이렇게 풀어내면 DP의 시간복잡도는 O(N)이 된다. 3겹의 for loop 중 첫번째 for문이 O(N), 내부 두 겹은 모두 O(1)이기 때문이다. j와 k는 N에 따라 변하는 값이 아니라는 점에서 O(1)임을 생각할 수 있다.

공간복잡도 또한 O(N)인데, 이는 dp 점화식 구조상 dp[0]과 dp[1]만 있어도 문제가 없음을 생각해볼 때 O(1)으로 최적화 가능하다. 나는 Tabulation이 눈에 잘 들어오는 풀이를 선호해서 O(N) 풀이 그대로 두었다.

 

Python 코드)

1) 백트래킹(Backtracking) - TLE

def backTrack(last_digit, digit_idx, used, val):
    if digit_idx == -1:
        if len(used) == 10 and N == len(str(val)):
            return 1
        else:
            return 0
    answer = 0
    if last_digit > 0:
        flag = False
        if last_digit - 1 not in used:
            used.add(last_digit - 1)
            flag = True
        answer += backTrack(last_digit - 1, digit_idx - 1, used, val * 10 + last_digit - 1)
        if flag:
            used.remove(last_digit - 1)
    if last_digit < 9:
        flag = False
        if last_digit + 1 not in used:
            used.add(last_digit + 1)
            flag = True
        answer += backTrack(last_digit + 1, digit_idx - 1, used, val * 10 + last_digit + 1)
        if flag:
            used.remove(last_digit + 1)
    
    return answer

N = int(input())

answer = 0
for i in range(11):
    answer += backTrack(i, N, set(), 0) % 1e+9

print(answer)

 

2) 동적계획법(Dynamic Programming) + 비트마스크(Bitmask) - AC

 

N = int(input())

MAX = (1<<10)

dp = [[[0] * MAX for _ in range(10)] for _ in range(N + 1)]
# dp[i][j][k]: number of stair numbers which has length i, last digit is j, and used bitmask is k.

for j in range(1, 10):
    dp[1][j][1 << j] = 1

for i in range(2, N + 1): 
    for j in range(0, 10):
        for k in range(MAX):
            if j > 0:
                dp[i][j - 1][k | (1 << (j - 1))] += dp[i - 1][j][k]
                dp[i][j - 1][k | (1 << (j - 1))] %= 1e+9
            if j < 9:
                dp[i][j + 1][k | (1 << (j + 1))] += dp[i - 1][j][k]
                dp[i][j + 1][k | (1 << (j + 1))] %= 1e+9
      
result = 0
for j in range(10):
    result += dp[N][j][MAX - 1]
    result %= 1e+9

print(round(result))