AtCoder Beginner Contest 198 D,F問題メモ

from itertools import permutations, starmap
from operator import mul

s1 = input()
s2 = input()
s3 = input()

whole = list(set(s1) | set(s2) | set(s3))
l = len(whole)
if l > 10:
    print('UNSOLVABLE')
    exit()

rev = {c: i for i, c in enumerate(whole)}

use = [0] * l
d = 1
for c in s1[::-1]:
    use[rev[c]] += d
    d *= 10
d = 1
for c in s2[::-1]:
    use[rev[c]] += d
    d *= 10
d = 1
for c in s3[::-1]:
    use[rev[c]] -= d
    d *= 10

ng_zero = {rev[s1[0]], rev[s2[0]], rev[s3[0]]}

for allocations in permutations(range(10), r=l):
    if any(allocations[i] == 0 for i in ng_zero):
        continue

    if sum(starmap(mul, zip(allocations, use))) == 0:
        t1, t2, t3 = s1, s2, s3
        for i in range(l):
            c = whole[i]
            d = str(allocations[i])
            t1 = t1.replace(c, d)
            t2 = t2.replace(c, d)
            t3 = t3.replace(c, d)
        print(t1)
        print(t2)
        print(t3)
        exit()

print('UNSOLVABLE')

import numpy as np
import numba


@numba.njit('i8[:,:](i8[:,:],i8[:,:],i8)', cache=True)
def mod_dot(a, b, MOD):
    """ 21x21行列の mod 10^9 あたりでのドット積はそのままやるとオーバーフローするので """
    n = len(a)
    c = np.zeros_like(a)
    for i in range(n):
        for j in range(n):
            c[i, j] = (a[i] * b[:, j] % MOD).sum() % MOD
    return c


def sum_patterns(s, coefs, MOD):
    """
    たとえば coefs=(3,2,5) のとき、3x+2y+5z=s となるような非負の(x,y,z)の組の個数
    coefsの合計値をkとして、O(k^3 log(s))
    """
    l = len(coefs)
    k = sum(coefs)
    arr = [0] * k
    dp = [0] * (1 << l)
    pp = [0] * (1 << l)
    for b in range(1, 1 << l):
        mlb = b & -b
        dp[b] = cur_dp = dp[b ^ mlb] + coefs[mlb.bit_length() - 1]
        pp[b] = cur_pp = pp[b ^ mlb] ^ 1
        if cur_pp == 0:
            arr[cur_dp - 1] -= 1
        else:
            arr[cur_dp - 1] += 1

    a = np.eye(k, k=-1, dtype=np.int64)
    a[0] = arr
    result = np.eye(k, dtype=np.int64)
    while s:
        if s & 1:
            result = mod_dot(result, a, MOD)
        a = mod_dot(a, a, MOD)
        s >>= 1

    return int(result[0, 0])


s = int(input())
MOD = 998244353

s0 = 30
s2 = 15
s2s2 = 8
s2s3 = 3
s2s4 = 2
s2s2s2 = 6
s3 = 5
s3s3 = 2
s4 = 2
s5 = 1
s6 = 1

patterns = [
    s0,
    s2,
    s2, s3,
    s2, s2s2, s3, s4,
    s2, s2s2, s2s2, s2s3, s3, s2s3, s4, s5,
    s2, s2s2, s2s2, s2s3, s2s2, s2s2s2, s2s3, s2s4,
    s3, s2s3, s2s3, s3s3, s4, s2s4, s5, s6
]

ans = 0

for b in range(32):
    coefs = []
    for i in range(5):
        if (1 << i) & b == 0:
            coefs.append(i + 1)
    coefs.append(6)

    sc = sum(coefs)
    if sc > s:
        continue

    ans = (ans + sum_patterns(s - sc, coefs, MOD) * patterns[b]) % MOD

print(ans)