AtCoder Beginner Contest 295 E,F,Ex問題メモ

AtCoder Beginner Contest 295

E - Kth Number

• $1～M$ からなる長さ $N$ の整数列 $A=(A_1,...,A_N)$ が与えられるが、いくつかは未定（$A_i=0$）である
• 以下の操作を1度ずつ順に行う
  • 各'0'を独立かつ等確率で $1～M$ に置き換える
  • $A$ を昇順に並び替える
• 操作後の $A_K$ の期待値を求めよ
• $1 \le K \le N \le 2000$
• $1 \le M \le 2000$

もとの $A$ の並び順は意味ないので、$0$ は抜き出して個数だけ数え、それ以外は昇順に並べておく。

$x=1～M$ のそれぞれについて、$A_K=x$ になる確率を求める。

x=7 K=10                            ※ 0は除いた上で昇順ソート済
i     1 2 3 4 5 6 7 8 9 10 ...         他に 0 が5個
A = ( _ _ _ 7 7 7 7 _ _ _  ... )       '_'は7以外の数字

$7$ が $K=10$ 番目に来るには、0から変換される数字について、以下の2つが両方成り立つ必要がある。

• $7$ 以下の数字が3個以上（これより少ないと、7が $i=10$ まで来ず、手前で終わってしまう）
• $7$ 未満の数字が7個未満（これより多いと、7未満の数字が $i=10$ まで来てしまう）

この3個、7個というのは、ソートされた $A$ 上で $7$ がどこからどこまで現れるかを $[L,R)$ として、 それぞれ　$\max(0,K-R+1), \max(0,K-L+1)$ で求められる。

元の $A$ には現れない数字でも、もし現れていればそこに来るだろう位置で $L=R$ として、同様に求められる。

で、この場合の数を求めたいが、直接、二項係数で計算するには 7以下の数字の個数と7未満の数字の個数、両方を固定せねばならず、$O(N^2)$ となってしまう。

ところが、逆に成り立たない条件を考えると、

• $7$ 以下の数字が3個未満、または
• $7$ 未満の数字が7個以上

この2つというのは、絶対に重複しない。（範囲の広いのが $a$ 個未満、範囲の狭いのが $b$ 個以上で、必ず $a \le b$ のため）

なので、独立に求められる。独立に求めるならそれぞれ $O(N)$ で可能。
それぞれ（7以下の数字の個数、7未満の数字の個数）を固定して求め、 全体 $M^{0の個数}$ から引いてやれば、$x=7$ の時の答えが求まる。

$\displaystyle \sum_{x} \frac{x \times (xの時の場合の数)}{M^{0の個数}}$ が答え。

計算量は $O(N(\log{N}+M))$

Python3

def precompute_factorials(n, MOD):
    f = 1
    factorials = [1]
    for m in range(1, n + 1):
        f = f * m % MOD
        factorials.append(f)
    f = pow(f, MOD - 2, MOD)
    finvs = [1] * (n + 1)
    finvs[n] = f
    for m in range(n, 1, -1):
        f = f * m % MOD
        finvs[m - 1] = f
    return factorials, finvs


n, m, k = map(int, input().split())
aaa = list(map(int, input().split()))
MOD = 998244353

facts, finvs = precompute_factorials(n, MOD)

free = aaa.count(0)
aaa = sorted(a for a in aaa if a > 0)
aaa.append(1 << 60)

k -= 1
l = 0
r = 0
ans = 0

all_pattern = pow(m, free, MOD)

for b in range(1, m + 1):
    while aaa[l] < b:
        l += 1
    while aaa[r] <= b:
        r += 1
    if l > k:
        break
    if r + free <= k:
        continue

    tmp_pattern = all_pattern
    for ika in range(max(0, k - r + 1)):
        pat = facts[free] * finvs[ika] * finvs[free - ika] % MOD
        pat *= pow(b, ika, MOD)
        pat %= MOD
        pat *= pow(m - b, free - ika, MOD)
        pat %= MOD
        tmp_pattern -= pat
        tmp_pattern %= MOD
    for miman in range(k - l + 1, free + 1):
        pat = facts[free] * finvs[miman] * finvs[free - miman] % MOD
        pat *= pow(b - 1, miman, MOD)
        pat %= MOD
        pat *= pow(m - b + 1, free - miman, MOD)
        pat %= MOD
        tmp_pattern -= pat
        tmp_pattern %= MOD

    ans += b * tmp_pattern
    ans %= MOD

ans *= pow(all_pattern, MOD - 2, MOD)
ans %= MOD
print(ans)

公式Editorialの方法。あまり見たことは無かったが、期待値を求める一つのテクニックっぽい。

• 適用条件: 正整数値を取る確率変数に対する期待値である
  • 取り得る値が正負入り交じったりしてると、下駄を履かせるなどして調整する必要がある

期待値は、確率変数を $X$ として次のように言い換えられる。

• $\displaystyle E = \sum_{x=1} (x \times (ちょうど X=x になる確率)) = \sum_{x=1} (x \le X になる確率)$

この変換をする場合は、仮に $X$ が偶数など特定の値しか取らないよ、とかいう場合でも、$x$ は $1～上限$ まで全ての値を合計する必要がある。

まぁ今回はそんなことは無く、$x=1～M$ まで $x \le X$ になる確率を求めればよい。

「ちょうど $X=x$」を求める場合には「$x-1$ 以下の個数」と「$x$ 以下の個数」を両方固定する必要があったが、 変換後なら「$x-1$ 以下の個数」さえ固定すれば求められるので、$O(N)$ でできる。

よくある言い換えとして、整数値（などの離散値）を取る確率変数の場合は、

• $X=5$ となる確率 $= (X \le 5$ となる確率 $)-(X \le 4$ となる確率 $)$

などのように、次に小さい値との差分とすることで「ちょうど」を「以下」に変換できるので、 この方法でも $x$ を固定した1回毎の計算を $O(N)$ にできる。

F - substr = S

• $f(x)$ を、'0'-'9'の文字列 $S$ により以下で定義する
  • $x$ を10進数で表記したとき $S$ が現れる回数（出現位置が重なっていてもよい）
  • 例) S=0202 のとき、f(10202)=1, f(3020202)=2
• $L$ 以上 $R$ 以下の全ての数 $k$ について、$f(k)$ の総和を $\mod{998244353}$ で求めよ
• 1つの入力につき最大1000ケース与えられる
• $1 \le |S| \le 16$
• $1 \le L \le R \lt 10^{16}$

公式解説は上手く言い換えて二分探索の問題に落とし込んでいるが、桁DPでも解ける。
ただ、きっちりと定義してやらないと遷移で混乱するし、公式の方が賢い。

$\sum_{k=L}^{R} f(k)=\sum_{k=1}^{R} f(k)-\sum_{k=1}^{L-1} f(k)$ なので、$[1,R]$ の範囲の問題が解ければよい。

• $DP[i][j][k] = $ 上から $i$ 桁目までを見て、既に先頭に0以外の数字が出てきていて、$S$ の $j$ 文字目までと一致していて、$k=0$ なら既に $R$ より小さいことが確定、$k=1$ なら $R$ ちょうどである数のパターン数

「$S$ の $j$ 文字目までと一致していて」の部分は、ちゃんというと、その数の $i-j+1,...,i-1, i$ 桁目の数字の並びと、$S$ の $1,2,...,j$ 文字目が一致しているということ。

$R$ の $i$ 桁目の数字を $c$、$S$ の $j$ 文字目を $d$ として、

• $DP[i-1][j-1][0]$ の次には必ず $d$ を置けるので、$DP[i][j][0]+=DP[i-1][j-1][0]$
• $DP[i-1][j-1][1]$ の場合、
  • $c=d$ なら同じ状態が継続、$DP[i][j][1] += DP[i-1][j-1][1]$
  • $c > d$ なら $d$ を置きつつ小さいことを確定できる $DP[i][j][0] += DP[i-1][j-1][1]$

これは、DP配列の $j=0$ に、上 $i$ 桁が $R$ より小さい/と一致 している数の個数を入れておくと、上記と一緒に考えられる。

$k=0$ の時は、$R$ を文字列としてみて先頭 $i-1$ 桁分より、1小さい数となる。

たとえば $R=12345,i=3$ なら、

• $k=0$ のとき、遷移元は 122 個
• $k=1$ のとき、遷移元は 1 個

$i>0$ の場合、$R$ より小さいことは確定している。

$S[0]=0$ なら、そこからの一致はできない。それ以外の場合は、1通りだけ存在する。

• $DP[i][1][0]+=1$（S[0]!=0 の場合）

$i=0$ の場合、既に何文字かの一致が見られる場合と似ているが、$S[0]=0$ の場合だけ注意して、

• $R[0]=S[0]>0$ なら、$k=1$ の状態で初手一致。$DP[0][1][1] += 1$
• $R[0]>S[0]>0$ なら、$k=0$ の状態で初手一致。$DP[0][1][0] += 1$

$DP[i][|S|]$ までたどり着いたら、出現したものが見つかったので、答えに加算する。

このとき、残り桁の自由度分だけ倍加して、

• $k=0$ の場合は残り $|R| - i$ 桁を自由に決められる。$10^{残り桁}$
• $k=1$ の場合は、$R$ の残りの下○桁を超えなければよいので、「$R$ の下○桁 + 1」（全て0の分）

をかけてから答えに加算する。

Python3

def solve(s, t):
    if t == '0':
        return 0

    ls = len(s)
    lt = len(t)
    if ls > lt:
        return 0

    r_int = int(t)
    s = list(map(int, s))
    t = list(map(int, t))

    res = 0
    dp_safe = [0] * (ls + 1)
    dp_just = [0] * (ls + 1)
    for i, c in enumerate(t):
        new_dp_safe = [0] * (ls + 1)
        new_dp_just = [0] * (ls + 1)

        new_dp_safe[0] += dp_safe[0] * 10
        new_dp_safe[0] += dp_just[0] * c
        new_dp_just[0] += dp_just[0]

        for j, d in enumerate(s, start=1):
            if c > d:
                new_dp_safe[j] += dp_just[j - 1]
            elif c == d:
                new_dp_just[j] += dp_just[j - 1]
            new_dp_safe[j] += dp_safe[j - 1]

        if i == 0:
            # c >= 1
            if c == s[0]:
                new_dp_just[1] += 1
            elif c > s[0] > 0:
                new_dp_safe[1] += 1
            new_dp_just[0] += 1
            new_dp_safe[0] += c - 1
        else:
            new_dp_safe[0] += 9
            if s[0] > 0:
                new_dp_safe[1] += 1

        res += new_dp_safe[-1] * pow(10, lt - i - 1)
        res += new_dp_just[-1] * (r_int % pow(10, lt - i - 1) + 1)

        dp_safe = new_dp_safe
        dp_just = new_dp_just

    return res


def bruteforce(s, l, r):
    ans = 0
    for k in range(int(l), int(r) + 1):
        k = str(k)
        i = 0
        cur = 0
        try:
            while i < len(k):
                i = k.index(s, i) + 1
                cur += 1
        except:
            pass

        if cur > 0:
            print(k, cur)
            ans += cur

    return ans


t = int(input())
for _ in range(t):
    s, l, r = input().split()

    # print(bruteforce(s, l, r))

    l = str(int(l) - 1)
    # print(s, l, r, solve(s, r), solve(s, l))
    ans = solve(s, r) - solve(s, l)
    print(ans)

Ex - E or m

• 各マス 0 or 1 からなる $N \times M$ のグリッドが与えられる。いくつかは'?'で未定である
• '?'を0か1で埋める方法のうち、以下の操作によって作られうるものの個数を $\mod{998244353}$ で求めよ
• 操作
  • $N \times M$ のグリッドを用意し、全て0で初期化する
  • 各行、独立に、左端から好きな列までを'1'にする（'E' のイメージ）
  • 各列、独立に、上端から好きな行までを'1'にする（'m' のイメージ）
• $1 \le N,M \le 18$

高速ゼータ変換の応用。

高速ゼータ変換というと、ある集合ごとに決められた値 $f(S)$ があって、 各集合 $S$ について、それを部分集合として含む全て集合についての $f(T)$ の和を求めるやつ。

• （例）$Z(01011) = f(01011) + f(01111) + f(11011) + f(11111)$

この実装は、各bitについて順番に「現在着目中のbitが立っている集合 $S$ の値を、そのbitを倒した集合 $S \ XOR \ bit$ に足し込む」という操作を繰り返して、$O(N2^N)$ でできる。

以下のDPを考える。

• $DP[i][S]=i$ 行目まで見て、上から'1'が繋がっている列がbitフラグ $S$ である状態で条件を満たす盤面のパターン数

盤面の状態と $S$ の対応の例示が公式Editorialにある。解説 - AtCoder Beginner Contest 295

bitフラグを2進数として率直に表記すると右の方が小さい添字を示すのに対し、配列では左の方が小さい添字になるのでややこしい。ここでは以降、配列に合わせ、2進数っぽい表記でも左の方が小さい添字を表すものとする。

$i-1$ から $i$ についての遷移を考える。

ある $S$ に対し、左から $k-1$ 列目まで'1'を伸ばし、$k$ 列目には'0'を埋める遷移を考える。
$S[j]$ を、bitフラグ $S$ で $j$ 列目を示す箇所のbit値（0/1）とする。

j         0123456789..
S         110101110111
i行目row  ?1?????1??0?  （入力に与えられたi行目）

k=5 の時  111110______  という状態を固定しての遷移となる

この時、

• 可否条件
  • $j=0,...,k-1$ の中に $row[j]=0$ である列が含まれる場合、'1'を $k$ まで伸ばすことができないので無理
  • $row[k]=1$ の場合、そこに'0'を置けないので無理
  • $j=k+1,...,M-1$ の中に $row[j]=1$ かつ $S[j]=0$ である列が含まれる場合、無理

それ以外なら遷移可能で、可能なときは

• $j=k+1,...,M-1$ の中で、
  • $row[j]=1$ の列は、必ずそのまま遷移
  • $row[j]=0$ の列は、$S[j]=1$ であったとしても必ず0に遷移
  • $row[j]=?$ の列は、$S[j]=1$ なら0/1両方に遷移。

j         0123456789..
S         110101110111
i行目row  ?1?????1??0?

k=5 の時
遷移先 S  110100_10_0_  （_ はそれぞれ独立に0/1両方に遷移）

要は、「$S$ から、$k$ 列目および以降のrowで'0'になっている列のbitを倒した $S'$」が最大であり（'_'に全て'1'を置いた状態を示す）、 そこから'_'、つまり「$k+1$ 列目以降で、$S'[j]=1$ かつ $row[j]=?$ になっている列」はそのbitを個別に1→0にした状態に遷移できる。

            v  v v
DP[i][110100010000] ← DP[i-1][S]
DP[i][110100010001] ← DP[i-1][S]
DP[i][110100010100] ← DP[i-1][S]
DP[i][110100010101] ← DP[i-1][S]
DP[i][110100110000] ← DP[i-1][S]
DP[i][110100110001] ← DP[i-1][S]
DP[i][110100110100] ← DP[i-1][S]
DP[i][110100110101] ← DP[i-1][S]  （この遷移先が S'）

しかし、愚直に全遷移先に足し込もうとすると、そのような列の数は最大 $M$、つまり遷移先は $2^M$ になり得る。
これは $(k,S)$ を固定した場合での遷移なので、$i$ ごとに $O(M 2^{2M})$ かかり、TLEである。

ここで、「部分集合に足し込む」ことを効率化する高速ゼータ変換が使える。

各 $(k,S)$ に対して $S'$ を求めてまとめて、 最後に高速ゼータ変換をおこなえば、部分集合に上手いこと足し込まれた結果が求まりそうである。

しかしここで問題となるのが、$k$ ごとに遷移可能な部分集合が変わってくる点。
「$k+1$ 列目以降で」$S'[j]=1$ かつ $row[j]=?$ の列の部分集合に遷移するので、 $k$ 列目までにそのような列があっても遷移してはいけない。

j         0123456789..
S         110101110111
i行目row  ?1?????1??0?

k=5       110100_10_0_  '_'の部分集合に遷移させたい

しかし、全(k,S)をまとめて最後に1回の高速ゼータ変換をおこなうと

          _10_00_10_0_  k以下の列の部分集合にも遷移しちゃう

よって、高速ゼータ変換を段階的に行う。
$k=5$ のときの $S'$ を $DP[i]$ に反映するのは、「$k$ 以下のbitについて、1→0 の足し込みを完了した後」にする。

こうすると、もはや $S'$ の $k$ 以下の列については遷移が完了しているので行われず、$k+1$ 以降の箇所のみ、ゼータ変換が適用される。

• $i=0,1,...,N-1$ に対し
  • $k=0,1,...,M-1$ に対し
    • $row[k]=1$ なら遷移不可、足し込みも不要（skip）
    • $row[k]=?$ なら、$k$ bit目に対する高速ゼータ変換の1→0の足し込みを行う
    • $row[:k]$ に $0$ が未登場なら、そこから遷移できる可能性がある
      • $DP[i-1]$ に含まれる各 $S$ に対し、
        • $j=k+1$ 以降で、$row[j]=1$ かつ $S[j]=0$ の列があれば、その $S$ からは遷移不可（skip）
        • その他の場合、$S$ から、$k$ と $row[j]=0$ である $j$ 列目を倒したbitsetを $S'$ とし、
        • $DP[i][S']←DP[i-1][S]$ とする
  • 最後に、row に 0 が出現しない場合、左から全てを'1'で埋めることができ、その場合は単純に前の値が引き継がれる
    • 各 $S$ に対し、$DP[i][S]←DP[i-1][S]$ とする

計算量は $O(NM2^M)$ となる。

Python3

def solve(n, m, field):
    MOD = 998244353
    p2m = 1 << m

    # 「上から'1'を繋げることができるかの列ごとのbitフラグ」ごとの暫定行目までの盤面パターン数
    #   ※左からの'1'によってカバーできるので繋げても切れてもよい、という場合は繋げるものとする（重複除外）
    dp = [0] * p2m
    dp[-1] = 1

    for i in range(n):
        ndp = [0] * p2m

        row = field[i]
        pos0 = 0
        pos1 = 0
        posx = 0
        for j in range(m):
            if row[j] == '0':
                pos0 |= 1 << j
            elif row[j] == '1':
                pos1 |= 1 << j
            else:
                posx |= 1 << j

        for j in range(m):
            # 左からはj-1列目まで伸ばし、j列目は'0'とするパターン

            if row[j] == '1':
                continue

            bit = 1 << j
            if j > 0 and row[j] == '?':
                for bitset in range(p2m):
                    if bitset & bit:
                        ndp[bitset ^ bit] += ndp[bitset]

            lo_mask = (1 << j) - 1
            if lo_mask & pos0:
                continue

            for bitset in range(p2m):
                if ~bitset & ~lo_mask & pos1:
                    continue
                ndp[bitset & ~bit & ~pos0] += dp[bitset]

        if pos0 == 0:
            for bitset in range(p2m):
                ndp[bitset] += dp[bitset]
                ndp[bitset] %= MOD
        else:
            for bitset in range(p2m):
                ndp[bitset] %= MOD

        dp = ndp

    return sum(dp) % MOD


n, m = map(int, input().split())
field = list(input() for _ in range(n))

ans = solve(n, m, field)
print(ans)