桁 DP の練習

ketadp.py

"""桁 DP の練習"""


def count(N):
    """N 以下の非負整数の数
    
    つまり、ただの N + 1 を動的計画法で回りくどく算出している。
    """
    zero = ord(b'0')

    keta = len(N)

    # dp[桁数][smaller]
    dp = [[0 for _ in range(2)] for _ in range(keta+1)]
    dp[0][0] = 1
    for i, _digit in enumerate(N):
        digit = _digit - zero
        for j in range(2):
            for d in range(10 if j else digit + 1):
                # smaller True からは True だけに遷移する。 N 以下なのは確定している。
                # smaller False のとき。
                # 見ている値 digit より小さければ遷移先は True 。
                # 同じであれば遷移先は False 。
                # 大きい場合の遷移先はなし。ループ変数 d に現れることがないようにする。
                dp[i + 1][j or d < digit] += dp[i][j]

    ans = dp[keta][0] + dp[keta][1]
    return ans


def mod(N, M):
    """N 以下で M で割り切れる非負整数の数
    
    つまり (int(N) // M) + 1 を動的計画法で算出している。
    """
    zero = ord(b'0')

    keta = len(N)

    dp = [[[0 for _ in range(M)] for _ in range(2)] for _ in range(keta+1)]
    # 解きたい問題のための情報を追加する。
    # dp[桁数][smaller][M で割った時のあまり] としてみる。
    dp[0][0][0] = 1
    for i, _digit in enumerate(N):
        digit = _digit - zero
        for j in range(2):
            for d in range(10 if j else digit + 1):
                for m in range(M):
                    dp[i + 1][j or d < digit][(m * 10 + d) % M] += dp[i][j][m]

    ans = dp[keta][0][0] + dp[keta][1][0]
    return ans


def main():
    N = b'1234'
    print(count(N))
    print(mod(N, 3))


if __name__ == '__main__':
    main()