AcWing 数位DP相关问题 1081. 度的数量

import sys
sys.stdin = open('data.txt', 'r')


sys.setrecursionlimit(999999)


from functools import lru_cache
@lru_cache(typed=False, maxsize=128000000)
def conbination_num(n, m):
    if m == n:
        return 1

    if m == 1:
        return n

    return conbination_num(n-1, m-1) + conbination_num(n-1, m)

# 小于等于N的所有可能的正整数个数
def get_num(n, k, b):
    arr = []
    val = n
    while val:
        arr.append(val % b)
        val //= b

    # 在更高位全部选择了最大数值状态为all_higher_maxd的情况下
    # 序列arr的前i位选择j位1，其余位为0的所有可能选择的种数
    @lru_cache(typed=False, maxsize=128000000)
    def dp(i, j, all_higher_max):
        if j > i+1:
            return 0

        if j == 0:
            return 1

        if i == 0:
            if j == 1:
                return 0 if all_higher_max and arr[0] == 0 else 1
            else:
                return 1
        else:
            if all_higher_max:
                if arr[i] == 0:
                    return dp(i-1, j, True)
                else:
                    return dp(i-1, j, False) + dp(i-1, j-1, arr[i] == 1)
            else:
                return conbination_num(i+1, j)

    return dp(len(arr)-1, k, True)

X, Y = map(int, input().split())
str = input()
try:
    K = int(str)
    B = int(input())
except:
    K, B = map(int, str.split())


print( get_num(Y, K, B) - get_num(X-1, K, B))