ACWing 数位DP相关问题 1086. 恨7不成妻

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

MOD = 1000000007
# 小于等于n的所有数字的平方和
def get_sum(n):
    if n == 0:
        return 0, 0, 0

    arr = []
    val = n
    while val:
        arr.append(val % 10)
        val //= 10

    # 前i个位置做选择，选择出的数字要求数位总和模7不等于bit_sum_mod7
    # 且数字本身模7不等于val_sum_mod7, 在所有高位都选择了最大值的标记为all_higher_max
    # 的情况下，所有合法数字的平方和，平方和是可以拆解成每个子集的数的平方和来求解的
    # 函数需要返回集合中所有数的数量，一次项累加和以及二次项累加和
    from functools import lru_cache
    @lru_cache(typed=False, maxsize=128000000)
    def dp(i, bit_sum_mod7, val_sum_mod7, all_higher_max):
        if i == 0:
            upper_bound = arr[i] if all_higher_max else 9

            cnt = 0
            sum = 0
            square_sum = 0
            for cur_val in range(upper_bound+1):
                if cur_val == 7:
                    continue

                if cur_val % 7 == bit_sum_mod7:
                    continue
                if cur_val % 7 == val_sum_mod7:
                    continue

                cnt += 1
                sum += cur_val
                square_sum += cur_val * cur_val

            return cnt, sum % MOD, square_sum % MOD

        else:
            upper_bound = arr[i] if all_higher_max else 9

            cnt = 0
            sum = 0
            square_sum = 0
            for cur_val in range(upper_bound + 1):
                if cur_val == 7:
                    continue

                new_cnt, new_sum, new_square_sum = dp(i-1, (bit_sum_mod7 - cur_val) % 7, (val_sum_mod7 - cur_val * 10 ** (i)) % 7, all_higher_max and cur_val == arr[i])
                cnt += new_cnt
                sum += new_cnt*(cur_val * (10**(i))) + new_sum
                square_sum += new_cnt * ( (cur_val**2) * (10**(2*i)) ) + 2*cur_val*(10**(i))*new_sum + new_square_sum

            return cnt, sum % MOD, square_sum % MOD

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


T = int(input())
for _ in range(T):
    a, b = map(int, input().split())
    print( (get_sum(b)[2] - get_sum(a-1)[2]) % MOD )