AtCoder Beginner Contest 198 D,F問題メモ

AtCoder Beginner Contest 198

D - Send More Money

問題

英大文字からなる3つの文字列 $S_1,S_2,S_3$ が与えられる
$S_1+S_2=S_3$ となる覆面算として解けるかどうか判定し、解ける場合は解の1つを求めよ
$1 \le |S_1|,|S_2|,|S_3| \le 10$
制限時間 $5$ sec.

解法

いかにもソルバがどっかに転がってそうではあるけど、まぁ総当たりすればよいので自力で書くことにする。

最大 $10! \times 10$ くらいの計算量が必要になるので、直接、文字をreplaceするなど下手な実装するとTLEになり得る。

文字毎に寄与をまとめておくとよい。

   SEND    E: S1の100の位, S2の1の位、S3の10の位に出現
+  MORE       → 100 + 1 - 10 = 91
= MONEY    E にたとえば6を当てはめると、91 x 6 = 546 の寄与

で、全文字の寄与の合計が0になると、その割り当ては覆面算の解ということになる。（先頭文字が0でないチェックも必要）

Python3

F - Cube

問題

6面に1つずつ正整数の書かれたサイコロを考える
6面の値の総和が $S$ になるようなサイコロへの数字の書き込み方の数を $\mod{998244353}$ で求めよ
回転して同じになるものは1つと数える（数字の向きは関係ない）
$6 \le S \le 10^{18}$

解法

$k$ 項間線形漸化式の $S$ 項目は高速に計算する方法があるよ、という問題。

正直、それを自力で導くのは相当骨が折れるので、知識問題感が強い。
（必ずしも使わなくても、パターン分類と行列累乗で解ける。公式解説の解法2。それでも数え上げに上手な発想の転換が必要だが）

以下のあたりが、高速な実装について触れられていて、簡潔と思った。Bostan–Mori のアルゴリズムというらしい。

線形漸化式を満たす数列の N 項目を計算するアルゴリズム
上記の具体例を交えた説明
- 線形漸化的数列のN項目の計算 - Qiita

Python3

        
          
          
              
              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)

            

        
      