HHKB Programming Contest 2022（AtCoder Beginner Contest 235）F,G問題メモ

HHKB Programming Contest 2022（AtCoder Beginner Contest 235）

中盤、妙に制約が厳しめの問題が多かった。

F - Variety of Digits

問題

$1$ 以上 $N$ 以下の整数で、以下の条件を満たすものの総和を $998244353$ で割ったあまりで求めよ
- 先頭に0をつけない10進数で表記したとき、$M$ 個の数字 $C_1,C_2,...,C_M$ がすべて出現する
$1 \le N \le 10^{10000}$
$1 \le M \le 10$
$0 \le C_i \le 9$

解法

制約や問題文が見るからに桁DP。あと使った数字の集合も持ちたいのでbitDPも絡んでくる。

桁が $\log_{10}N \le 10^4$ 個、1桁あたりの遷移に $M \times 2^M \lesssim 10^4$ かかるので、ちょっときつくないか？という感じはするものの、他に方法も思いつかない。
まぁ遷移はシンプルなので下手な書き方さえしなければいけると信じよう。

桁DPは個数を数えさせるものをよく見るが、今回は総和であることに注意。

以下のDPを定義する。
公式解説では、使った数字の集合を $C_i$ に関係なく $0～9$ の全てで区別しているが、ここでは $C_i$ に含まれる数字のみで区別している。

$DP_{count}[i][S]=N$ の先頭 $i$ 桁だけで考えたとき、使った $C_i$ の集合が $S$ であり、$N$ より小さいことが既に確定している数字の個数
$DP_{total}[i][S]=$ 同、総和

$S$ をbit集合で表現すると、答えは $DP_{total}[N][2^M-1]$ となる。
ただし $N$ 自身も条件を満たす場合、それは含まれていないので最後に足す。

桁DPは、区別すべき状態が $S$ の他にも暗黙的に存在するため、きちんと分けて遷移を考える必要がある。

①先頭 $i$ 桁までで、$0$ 以外の数字が既に現れていて、$N$ より小さいことが確定している
- 上記のDPに数えられているのはこれのみ
②先頭 $i$ 桁までで、まだ $0$（leading-zero）以外の数字が現れていない
③先頭 $i$ 桁まで $N$ と一致している

②は、問題によっては①と同等に扱ってよいこともあるが、今回は $C_i=0$ があった場合に先頭に出てくるか途中に出てくるかで意味合いが異なるため、分けて考える必要がある。

②③に該当する状態も、ある桁で①に変化すればそこからDPに含めてやる必要がある。
その際、特に③で必要になるため、以下の情報も合わせて管理する。

$just_{status} =$ $N$ の先頭 $i$ 桁までに出現した $C_i$ の集合
$just_{total} =$ $N$ の先頭 $i$ 桁 $\mod{998244353}$

初期化

1桁目に関しては、下記の遷移のどれとも異なるため、個別に行う。1桁目の数字を $d$ とすると、

0を置いた場合②、$d$ を置いた場合③の状態になり、$1～d-1$ を置いた場合に①としてDPの管理対象になる。

更新できる数字の範囲が違うだけで、式は②と同様となる。

①の遷移

$i+1$ 桁目に $0～9$ の何を置いてもよい（$N$ 以下という条件が破られることはない）。

数字 $d$ を置くとして、以下のような遷移となる。遷移先のみが異なる。

$d$ が $C_k$ に該当する場合
- 使った数字の集合は $S'=S|(1<<k)$ に更新される
- $DP_{count}[i+1][S']+=DP_{count}[i][S]$
- $DP_{total}[i+1][S']+=DP_{total}[i][S] \times 10 + DP_{count}[i][S] \times d$
$d$ が $C$ に含まれない場合
- $DP_{count}[i+1][S]+=DP_{count}[i][S]$
- $DP_{total}[i+1][S]+=DP_{total}[i][S] \times 10 + DP_{count}[i][S] \times d$

$N=30000, C=(1,3,5)$ として、$DP_{total}[3][0b011]$ には $103, 113, 31$ などが該当し、その総和が記録されている。

$4$ 桁目に $5$ を付け足す場合、$1035, 1135, 315$ などの総和を $DP_{total}[4][0b111]$ に加えたい。

それには、$DP_{total}[3][0b011]$ を10倍した上で、個数分だけ $5$ を加えればよいので、上のような遷移となる。

②の遷移

この状態に該当するパターン数は常に1通り、総和は0。

$i+1$ 桁目にも0を置いた場合はこの状態が継続される。

2桁目以降、$1～9$ を置いた場合はDPの管理対象となる。
①と同じような遷移となるが、パターン数と総和が決まっているためシンプルになる。

$d$ が $C_k$ に該当する場合
- $DP_{count}[i+1][1<<k]+=1$
- $DP_{total}[i+1][1<<k]+=d$
$d$ が $C$ に含まれない場合
- $DP_{count}[i+1][0]+=1$
- $DP_{total}[i+1][0]+=d$

③の遷移

この状態に該当するパターン数も常に1通り、総和は $just_{total}$。

$N$ の $i+1$ 桁目を $g$ とすると、$g$ を置いたらこの状態が継続される。

$g$ を超える数は $N$ を超えるので論外。

$0～g-1$ の数を置くと、$N$ 未満になることが確定し、DPに含められる。$d$ を置くとして、

$d$ が $C_k$ に該当する場合
- 使った数字の集合は $S'←just_{status}|(1<<k)$ に更新される
- $DP_{count}[i+1][S']+=1$
- $DP_{total}[i+1][S']+=just_{total} \times 10 + d$
$d$ が $C$ に含まれない場合
- $DP_{count}[i+1][S]+=1$
- $DP_{total}[i+1][S]+=just_{total} \times 10 + d$

Python3

import os
import sys

import numpy as np


def solve(n, m, ccc):
    MOD = 998244353

    lim = 1 << m
    bits = np.zeros(10, np.int64)
    for i, c in enumerate(ccc):
        bits[c] = 1 << i

    dp_from = np.zeros((lim, 2), np.int64)
    dp_to = np.zeros((lim, 2), np.int64)

    others = np.zeros((2, 11, 2), np.int64)  # [i,j) のうち、cccに含まれないものの(個数,総和)
    others[0, 1, 0] = int(bits[0] == 0)
    for i in range(1, 10):
        if bits[i] != 0:
            others[:, i + 1, :] = others[:, i, :]
        else:
            others[:, i + 1, 0] = others[:, i, 0] + 1
            others[:, i + 1, 1] = others[:, i, 1] + i

    d = n[0]
    for c in ccc:
        if c == 0:
            continue
        if d <= c:
            break
        dp_from[bits[c], 0] = 1
        dp_from[bits[c], 1] = c
    dp_from[0] = others[1, d]

    just_status = bits[d]
    just_total = d

    for d in n[1:]:

        for i in range(lim):
            for c in ccc:
                ni = i | bits[c]
                dp_to[ni, 0] += dp_from[i, 0]
                dp_to[ni, 1] += dp_from[i, 1] * 10 + c * dp_from[i, 0]
            dp_to[i, 0] += dp_from[i, 0] * others[0, 10, 0]
            dp_to[i, 1] += dp_from[i, 1] * 10 * others[0, 10, 0] + others[0, 10, 1] * dp_from[i, 0]

        for c in ccc:
            if d <= c:
                break
            ni = just_status | bits[c]
            dp_to[ni, 0] += 1
            dp_to[ni, 1] += just_total * 10 + c
        dp_to[just_status, 0] += others[0, d, 0]
        dp_to[just_status, 1] += just_total * 10 * others[0, d, 0] + others[0, d, 1]

        for c in ccc:
            if c == 0:
                continue
            ni = bits[c]
            dp_to[ni, 0] += 1
            dp_to[ni, 1] += c
        dp_to[0] += others[1, 10]

        dp_to %= MOD
        dp_to %= MOD

        just_status |= bits[d]
        just_total = (just_total * 10 + d) % MOD

        dp_from, dp_to = dp_to, dp_from
        dp_to.fill(0)

    ans = dp_from[-1, 1]
    if just_status == lim - 1:
        ans += just_total
    ans %= MOD

    return ans


SIGNATURE = '(i8[:],i8,i8[:])'
if sys.argv[-1] == 'ONLINE_JUDGE':
    from numba.pycc import CC

    cc = CC('my_module')
    cc.export('solve', SIGNATURE)(solve)
    cc.compile()
    exit()

if os.name == 'posix':
    # noinspection PyUnresolvedReferences
    from my_module import solve
else:
    from numba import njit

    solve = njit(SIGNATURE, cache=True)(solve)
    print('compiled', file=sys.stderr)

n = np.fromiter(map(int, input()), dtype=np.int64)
m = np.int64(input())
ccc = np.fromiter(map(int, input().split()), dtype=np.int64)
ans = solve(n, m, ccc)
print(ans)

G - Gardens

問題

$N$ 個の庭と、3種類の作物の苗（リンゴ、バナナ、チェリー）がそれぞれ $A,B,C$ 株ある
- 庭は区別できる。同一種類の苗は区別しない
以下の条件を満たして苗を植える方法の個数を $998244353$ で割ったあまりで求めよ
条件
- 各苗は1つの庭に「植えない」「1株植える」のどちらか（同一種類を2株以上植えない）
- どの作物も植わっていない庭があってはいけない
- 各苗は余ってもよい
$1 \le N \le 5 \times 10^6$
$0 \le A,B,C \le N$

解法

仮に、何も植わっていない庭を許すとするならば、各作物独立に考えてよくなるので簡単になる。

$f(n,a)=n$ 個の庭に、1種類の作物を $a$ 株以下植える方法の個数
$={}_nC_a+{}_nC_{a-1}+...+{}_nC_0$

とすると、$f(N,A) \times f(N,B) \times f(N,C)$ で求められる。

こういう場合は包除原理が使えそう。つまり、

    少なくとも0個の庭に何も植わっていない場合数
  - 少なくとも1個の庭に何も植わっていない場合数
  + 少なくとも2個の庭に何も植わっていない場合数
    ...
+/- 少なくともN個の庭に何も植わっていない場合数
-----------------------------------------------
  = 答え

少なくとも $k$ 個の庭に何も植わっていない場合数は、以下をかけあわせた数となる。

何も植えないと決める庭の選び方 ${}_NC_k$
その他の庭への植え方 $f(N-k,A) \times f(N-k,B) \times f(N-k,C)$

ただし $n \lt r$ の場合、${}_nC_r=0$ とする。

上記の数え方だと、例えば以下の2つは結果的に同じなのに別々に数えられている。

庭1に何も植えないと決め、結果的に庭2にも何も植わらなかった
庭2に何も植えないと決め、結果的に庭1にも何も植わらなかった（他の庭の状態も同じ）

一般に、数え上げでは重複に気をつけることが多いので、それっていいの？と躊躇してしまうが、包除原理の場合は「確実に植わっていない特定の $k$ 個」が異なっていれば後から適切に除かれるのでこれでいい。

「植えないと決める対象とする/しない」を各庭に対して決めた $2^N$ 通りについて、対象数が偶数なら正、奇数なら負で場合数を足し合わせるという手法が本来の包除原理であり、今回は対象数が同じもの同士をまとめられるのでまとめている、と考えられる。

さて、これで $k=0～N$ について $f(N-k,A) \times f(N-k,B) \times f(N-k,C)$ を求めればよいとわかった。

しかし $f$ の内部では二項係数を $0～A$ まで足し合わせているため、$k$ をあわせるとまだ $O(N^2)$ かかる。

二項係数の途中までの総和

$f(n,A)$ と $f(n+1,A)$ の差分に着目する。

$f(N,A)$ が表しているのは、パスカルの三角形の $N$ 段目の、左から第 $A$ 項までの和である。（0-index）

          1
        1   1
      1   2   1
    1   3   3   1
  1   4   6   4   1
1   5  10  10   5   1

f(5, 3) = 1+5+10+10 = 26
f(4, 3) = 1+4+6+4   = 15
f(3, 3) = 1+3+3+1   =  8
f(2, 3) = 1+2+1     =  4
f(1, 3) = 1+1       =  2
f(0, 3) = 1         =  1

パスカルの三角形は、自身の斜め上の2つの数字の和が、自身の値となる。

そのため、$f(n,a)$ と $f(n+1,a)$ を比較した場合、 $n$ 段目の数字はほとんど $n+1$ 段目には左右に2回ずつ波及している。

唯一、一番右端の第 $a$ 項のみ、左にしか波及しない。

  1   4   6   4
↙↘↙↘↙↘↙ ×
1   5  10  10   -

よって、まず $f(0,A)=1$ からスタートし、 $f(n+1,A)=2f(n,A)-{}_nC_{A}$ として逐次的に計算していくと、 $O(N)$ で $f(0,A)～f(N,A)$ を全て求められる。

二項係数周りで「$O(N)$ を $O(1)$ にする」便利な公式は多数あるけど、こういう「$O(N^2)$ を $O(N)$ にする」テクニックは検索に頼りづらい。