문제 링크

<aside> 💡 다이나믹 프로그래밍/그리디 알고리즘

</aside>

Memory 170488KB Time 488ms Code Length 730B

MOD = int(1e9 + 7)

def DP(a, b, c, d, memo):
    if a == n:
        return int(not b)
    if (a, b, c, d) in memo:
        return memo[(a, b, c, d)]

    ret = 0
    if d < 10:
        for i in range(c, 10):
            if c + d == i:
                ret += DP(a + 1, b, c + d, d, memo)
            else:
                ret += DP(a + 1, b - 1, i, 10, memo)
            ret %= MOD
    else:
        for i in range(c, 10):
            ret += DP(a + 1, b, i, i - c, memo)
            ret %= MOD

    memo[(a, b, c, d)] = ret
    return ret

n, k = map(int, input().split())

if k > 9:
    print(0)
else:
    ans = 0
    memo = {}
    for i in range(1, 10):
        ans += DP(1, k - 1, i, 10, memo)
        ans %= MOD
    print(ans)

이 코드는 주어진 조건에 따라 숫자를 생성하고 그 숫자의 개수를 세는 것이다.

먼저, MOD는 나머지 연산을 위한 상수이다.

DP 함수는 동적 프로그래밍 알고리즘을 구현한 것으로, a는 현재까지 생성된 숫자의 길이, b는 남은 숫자의 개수, c는 마지막으로 사용된 숫자, d는 마지막 두 숫자의 차이, memo는 중복 계산을 방지하기 위해 이전에 계산한 결과를 저장하는 딕셔너리이다.

만약 a가 n과 같다면, 모든 숫자를 사용한 것이므로 b가 0인지 확인하고 그 결과를 반환한다.

만약 (a, b, c, d)가 memo에 있다면, 이전에 계산한 결과를 반환한다.

그렇지 않다면, d가 10보다 작은 경우와 그렇지 않은 경우로 나눠서 계산한다.

d가 10보다 작은 경우, c부터 9까지의 숫자 i를 사용해 새로운 숫자를 생성하고, 그 때의 경우의 수를 ret에 더한다.

d가 10 이상인 경우, c부터 9까지의 숫자 i를 사용해 새로운 숫자를 생성하고, 그 때의 경우의 수를 ret에 더한다.

마지막으로, 계산한 결과를 memo에 저장하고 반환한다.

n과 k는 사용자로부터 입력받는다.

만약 k가 9보다 크다면, 조건에 맞는 숫자를 생성할 수 없으므로 0을 출력한다.

그렇지 않다면, 1부터 9까지의 숫자 i를 사용해 숫자를 생성하고, 그 때의 경우의 수를 ans에 더한다.

마지막으로, ans를 출력한다.