AtCoder Regular Contest 170 B,C,D,E問題メモ

AtCoder Regular Contest 170

ARC170で検索すると、スターウォーズに出てくる宇宙戦闘機が出てくる。

B - Arithmetic Progression Subsequence

• 各要素が $1～10$ の整数である長さ $N$ の数列 $A$ が与えられる
• $l \le r$ の整数組 $(l,r)$ で、以下の条件を満たすものの個数を求めよ
  • $(A_l,...,A_r)$ に含まれる（連続とは限らない）長さ3の部分列のうち、等差数列になっているものが存在する
• $3 \le N \le 10^5$

満たさないものを数える。

要素が1～10しかないので、ちょっと長くなれば何かしらの等差数列はできてしまいそう。
少なくとも同じ整数が3つ出てきたら公差0の等差数列なので、長さ（$r-l+1$）が21以上になることはない。

なので、全ての $(l, r)$ の組の個数 $\dfrac{N(N+1)}{2}$ をとりあえずの答えとして、

• $l=1,2,...$ を固定する
• $r=l,l+1,...$ を、等差数列ができるまで試す
• 等差数列ができたら、その直前までの組の個数 $r-l$ が、答えから除かれる

とすれば、残ったものが答え。

$l$ を固定した時の $r$ がそこまで長くならないので愚直に調べても大丈夫という、 計算量の適切な見積もり（または直感）が必要な問題。

あり得る最大の“等差数列ができない区間の長さ”を $k$ として、$O(N k^2)$ などで計算できる。

Python3

n = int(input())
aaa = list(map(int, input().split()))

ans = n * (n + 1) // 2
for l in range(n):
    ng = set()
    for r in range(l, n):
        b = aaa[r]
        if b in ng:
            break
        for k in range(l, r):
            a = aaa[k]
            c = b + b - a
            if 1 <= c <= 10:
                ng.add(c)
    else:
        r = n
    ans -= r - l

print(ans)

C - Prefix Mex Sequence

• 各要素が $0$ または $1$ からなる長さ $N$ の数列 $S$ が与えられる
• 以下の条件を満たす長さ $N$ の数列 $A$ の個数を $\mod{998244353}$ で求めよ
  • 各要素は $0$ 以上 $M$ 以下
  • $i=1,2,...,N$ について、
    • $S_i=1$ ならば、$A_i=mex(A_1,A_2,...,A_{i-1})$
    • $S_i=0$ ならば、$A_i \neq mex(A_1,A_2,...,A_{i-1})$
• $1 \le N \le 5000$
• $0 \le M \le 10^9$

$m(i)=mex(A_1,A_2,...,A_i)$ とする。

• $m(i)$ は $i$ が進んで減ることはなく、$S_i=1$ となる $i$ で少なくとも1増え、$S_i=0$ では増えない

S      0 1 0 0 1 0 1 1
Aの例  2 0 3 5 1 2 4 6
m(i)   0 1 1 1 4 4 6 7

$S_i=1$ の時、$A_i$ は $m(i-1)$ で一意に決まっている。
一方、$S_i=0$ の時、$m(i-1)$ 以外なら何を置いてもよい。

$S_i=0$ の時に $m(i)+1$ とか $m(i)+2$ とかが埋められてたら、 次に $S_i=1$ で $m(i)$ が埋められたとき、その次のMexは一度にたくさん増えることもある。

S      0 1 0 0 1
Aの例  2 0 3 5 1
m(i)   0 1 1 1 4
       ^   ^      ←次のMexの1より1～2多い 2,3がここで埋められているため、
               ^  ←その後に1が埋められると、次のMexは一気に4まで飛ぶ

この性質のために、DPをするにしても「現在のMex」をキーとしても、次のMexを特定できないためあまり上手くいかない。

もし、$N$ より $M$ の方が十分大きければ、

• $S_i=0$ のとき、$0～M$ の $M+1$ 個の中から、その時のMexに当たらない $M$ 通りの数の中から選んで置く
• $S_i=1$ のとき、（考えられる $A$ は様々あるにしろ、それぞれの $A$ につき）$m(i-1)$ は一意に決まっているので、1通り

なので、$M^{Sに出てくる0の個数}$ が答えとなる。

その方法で上手くいかないのは、$M$ が小さくて、$S_i=1$ でMexを更新したいのに次のMexが $M+1$ 以上となってしまう場合。

S      0 1 0 0 1 0 1 1
Aの例  2 0 3 5 1 2 4 x   ←M=5 のとき、6を置けなくて破綻
m(i)   0 1 1 1 4 4 6

破綻するのは、それまでに $0～M$ が全て出現しているのに、$S_i=1$ が来てしまった場合。
つまり、「$0～M$ のうち、出現した種類数」を管理することで、破綻するかどうかを判別可能となる。

• $DP[i,j]=A_i$ までを決め、既に使われた種類数が $j$ 個の時のパターン数

• $S_{i+1}=0$ のとき、
  • 既に過去に使われた数を置く $DP[i+1,j]+=DP[i,j] \times j$
  • $m(i)$ 以外の新しい数を置く $DP[i+1,j+1]+=DP[i,j] \times (M-j)$
• $S_{i+1}=1$ のとき、
  • 必ず新しい数 $m(i)$ を置く $DP[i+1,j+1]+=DP[i,j]$

で、$DP[N,1～M]$ の和が答え。$M$ 自体はでかい可能性があるが、 実質取り得る値は最大 $N$ なので、DP配列は $N \times \min(N,M)$ だけ用意しておけばよい。

Python3

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

dp = [1]
for a in aaa:
    ndp = [0] * (len(dp) + 1)
    if a == 0:
        for j in range(len(dp)):
            ndp[j] += dp[j] * j
            ndp[j + 1] += dp[j] * (m - j)
    else:
        for j in range(len(dp)):
            ndp[j + 1] += dp[j]

    for j in range(len(ndp)):
        ndp[j] %= MOD
    dp = ndp

ans = sum(dp[:m + 2])
ans %= MOD
print(ans)

D - Triangle Card Game

• AliceとBobがゲームをする
• 2人にはそれぞれ $N$ 枚のカードが配られている
  • Aliceのカードには $A=(A_1,...,A_N)$ の正整数が書かれている
  • Bobのカードには $B=(B_1,...,B_N)$ の正整数が書かれている
• ゲームは以下の手順で行われる
  • Aliceが1枚、手持ちのカードを選んで出す
  • Bobが1枚、手持ちのカードを選んで出す
  • Aliceが1枚、（1手目に出した以外の）手持ちのカードを選んで出す
  • 3枚のカードに書かれた値を3辺の長さとする（非退化な）三角形が作れればAliceの勝ち、それ以外はBobの勝ち
• どちらが必勝か、$T$ 個のテストケースそれぞれに付き判定せよ
• $1 \le T \le 10^5$
• $1 \le A_i,B_i \le 10^9$
• 1つの入力に含まれる $N$ の総和 $\le 2 \times 10^5$

使う知識は「三角形が作れる ⇔ 3辺の長さ $a,b,c$ で一番長い辺を $c$ として、$a+b \gt c$」のみ。
だが、細かな場合分けがなかなか難しい。

無理なケース、可能なケースを洗い出してみて、なるべく簡潔に全てのパターンを網羅できる考え方にまとめあげたい。

Alice 1手目, Bob, Alice 2手目が出すカードの整数を、それぞれ $a,b,c$ とする。

まず $a$ を固定して全通り試す。
各 $a$ に対し $b$ を、$a$ より小さい場合と大きい場合で分け、どのような $c$ に対しても勝てる $b$ があるか調べる。

最長辺の候補は、$a$ か $c$。

三角形を作れなくなるのは以下の2通り。

|------a------|   a >= b+c1
|--b--||--c1-|

|------a------||--b--|   a+b <= c2
|----------c2----------|

Bobにとって、$b$ は $B$ の中で最も小さい数 $B_{min}$ を使うとしてよい。
Aliceにとって、$c_1$ は $a$ の次に大きい数、$c_2$ は $a$ の次に小さい数を使うことにしてよい。

$b,c$ ともに $a$ に対する最適なカード候補が確定するので、実際に確認する。

最長辺の候補は、$b$ か $c$。

三角形を作れなくなるのは以下の2通り。

①
|--a--||--c1--|      a+c1 >= b
|-------b-------|

②
|--a--||------b------|   a+b <= c2
|----------c2----------|

こちらは、$a \gt b$ の場合と比較して少々難しい。
$B_{min}$ のような、「$b,c$ はこれを選んどきゃいい」みたいな決定版が無い。

$c_1,c_2$ を、$A$ の大きさ順で隣り合う数とする。
この間の“ちょうどいい” $b$ を出されたとき、 「$c_1$ 以下を出したら①になるし、$c_2$ 以上を出したら②になる」という事態に陥り、Aliceは負ける。

ちょうどいい $b$ の範囲とは、$c_1+a \le b \le c_2-a$ となる。
つまり、$c_1,c_2$ を互いに $a$ ずつ縮めた範囲となる。

 c1                c2
  |--a-->    <--a--|
        |----|         ←この間にbが存在したらBobの勝ち
 
 c1=Amax
  |--a-->
        |---------...  ←または、c1がAmaxの場合、こうなる

いずれにしても、「Aliceは $X$ 以下の $a$ を出したら負ける」という形の制約となる。
各 $b$ につき、$A$ の中で $b$ に最も近い両側の $c_1,c_2$ を調べれば、$X$ が分かる。

では、$X$ より大きな $a$ を出せば $a \le b$ の条件下ではAliceは負けることが無い……かというと、そうではない。

$c_1,c_2$ と $a$ が、同じカードとなる場合がある。

 c0         c1=a             c2
             |--a-->    <--a--|
                   |----|         ←この間にbは存在しなくても、、
                   
  |--a-->               <--a--|     c1をaとして使ってしまった場合、
        |---------------|         ←bが存在したらAliceが負ける範囲はこっちになる

これは、$a$ を全探索する時に両隣の数を確認し、$c_0+a,c_2-a$ の間に $b$ が存在しないかを確認すればよい。

まとめると、

• 「Aliceがこれ以下を初手で出したら負ける値 $X$」を求める
• $X$ より大きい $a$ を全て試す
  • $A$ 内の大きさ順で $a$ の両端の数を $c_0,c_2$ とする
  • $c_0+a \le b \le c_2-a$ なる $b$ が存在したらダメ（$a \le b$ の場合の条件）
  • そうでない場合、$a \ge B_{min} + c_0$ かつ $a+B_{min} \le c_2$ だったらダメ（$a \gt b$ の場合の条件）
  • そうでない $a$ が1つでも見つかれば、Aliceの勝ち
• Aliceが勝てる $a$ が存在しなければ、Bobの勝ち

Python3

from bisect import bisect_right, bisect_left


def solve(n, aaa, bbb):
    aaa.sort()
    bbb.sort()
    b_min = bbb[0]

    a_lose = max(0, bbb[-1] - aaa[-1])
    for i in range(n - 1):
        a1 = aaa[i]
        a2 = aaa[i + 1]
        if a1 == a2:
            continue
        l = bisect_right(bbb, a1)
        r = bisect_left(bbb, a2)
        if l == n:
            break
        for j in range(l, r):
            b = bbb[j]
            a_lose = max(a_lose, min(b - a1, a2 - b))

    for i in range(n - 1):
        a1 = aaa[i]
        if a1 <= a_lose:
            continue
        a2 = aaa[i + 1]
        if i > 0:
            a0 = aaa[i - 1]
            l = bisect_left(bbb, a0 + a1)
            r = bisect_right(bbb, a2 - a1)
            if l < r:
                continue
            if a1 + b_min > a2 or a0 + b_min > a1:
                return 0
        else:
            if a1 + b_min > a2:
                return 0
    return 1


t = int(input())
for _ in range(t):
    n = int(input())
    aaa = list(map(int, input().split()))
    bbb = list(map(int, input().split()))

    res = solve(n, aaa, bbb)
    if res == 0:
        print('Alice')
    else:
        print('Bob')

E - BDFS

• $N$ 頂点 $N$ 辺のグラフがあり、頂点 $1,2,...,N,1$ の順に円状に繋がっている
• 頂点1から、$\frac{P}{100}$ で深さ優先探索、$1-\frac{P}{100}$ で幅優先探索を確率的に選択するような探索をする（詳細は問題ページ参照）
• 始点からの距離配列 $D=(D_1,...,D_N)$ について、探索終了後の $D$ の総和の期待値を $\mod{998244353}$ で求めよ
• 1つの入力につき $T$ ケース与えられる
• $1 \le T \le 10^4$
• $2 \le N \le 10^{18}$
• $1 \le P \le 99$

なんか問題文は複雑だが、大部分はBFSやDFSでいつもやっていることの説明に過ぎない。
「先頭に追加するか、末尾に追加するか確率的に選択される」という点が変わっている。

グラフが円状なので、$Q$ には終始、高々2要素しか入っていることはない。
結局、「時計回りルートと反時計回りルート、どちらの探索を進めるかが確率的に選択される」という問題になる。

直前に進めたルートと同じルートを確率 $p \ (=\frac{P}{100})$ で進め、反対のルートを $q \ (=1-p)$ で進めることになる。

2本のルートしかないので、$D_v$ が取り得る値は、 時計回りルートの $v-1$ か、反時計回りルートの $N-v+1$ しかない。

            Dvの値の一例
⑧--①--②  １--０--１
｜      ｜  ｜      ｜
⑦      ③  ２      ２
｜      ｜  ｜      ｜
⑥--⑤--④  ３--４--５

「グリッドを左上から右下まで最短距離で移動する経路の数え上げや確率に帰着させる」的な言い換えはよくあるが、 この問題もできるっちゃできる。（ただし、後述の通りTLEとなる）

今回の場合、下移動を時計回り、右移動を反時計回りに1進むことに対応させられる。

（合計 $N-1$ 回移動した箇所のマスに進む確率）×（それぞれの場合のDの合計）を、各マスについて合計すれば答えとなるのだが、 ただ、$N$ の上限が $10^{18}$ なので、そのままでは無理っぽい。
まぁ、ひょっとしたら上手くまとめられるかもしれないし、イメージをつかむつもりで進めてみる。

N=5
                時計回り   反時計回り  Dの合計
□□□□■  ←  0,1,2,3,4  0           10
□□□■    ←  0,1,2,3    0,1          7
□□■      ←  0,1,2      0,1,2        6
□■        ←  0,1        0,1,2,3      7
■          ←  0          0,1,2,3,4   10

ちょっと特殊なのが、右に進む、下に進むに対して進む確率が決まっているわけではなく、 「直前の方向を継承する場合に $p$、方向転換をする場合に $q$」の確率で選択されるということ。

なので、あまり一般的な経路数え上げのセオリーは使えない。

一応、以下のDPで計算できるし、

• $DP[i,j,k]=i$ 個下に、$j$ 個右に進み、直前の移動が $k=0:下,1:右$ であるような状態になる確率
  • $DP[0,0,1]=1$（右を向いた状態）からスタート

もう少し考えると、1から埋めていかなくても、 「どの列で方向転換するか」「方向転換した列で、それぞれ何回下がったか」により、組合せで求められはする。

• $\displaystyle DP[i,j,0] = \sum_{e=0}^{\min(i-1,j)} {}_jC_{e}{}_{e+1}H_{i-e-1}p^{N-2e-1}q^{2e+1}$
• $\displaystyle DP[i,j,1] = \sum_{e=1}^{\min(i,j)} {}_jC_{e}{}_{e}H_{i-e}p^{N-2e}q^{2e}$

ただ、$e$ のループで $O(N)$ かかるし、仮にこれが求められたとして、1頂点ずつ求めていったらさらに $O(N)$ かかる。
まとめて綺麗な形にする方法もパッと見えない。

別のアプローチが必要となる。

前述の通り、「時計回りルートと反時計回りルートの移動回数が $(x,y)$ となる確率」を求めて、 そこに「その時の $D$ の和」をかける方針だとTLEなので、 どうにかして「移動回数が $(x,y)$ となる確率 × その時の $D$ の和」の $x=1～N-1$ の総和を直接的に求める、ということをしないといけない。

ここで解説を読むが、まぁ、これは思いつかないわ、という感じ。

• 下に $L$ 回、右に $R$ 回（$L+R=N-1$）移動するような経路を1つ取り出すと、
  • この時の $D$ の総和は、$\dfrac{L(L+1)}{2}+\dfrac{R(R+1)}{2}$ で決まっている
• ここで、$\dfrac{L(L+1)}{2}$ というのは「$L$ 個のものの中から重複を許して2個選ぶ組合せの数 ${}_{L}H_{2}$」と一致する¹⁾
• なのでDPで「下移動を暫定いくつ選んだか 0～2」を状態に加えれば、2個選んだ状態の値が「そのマスへの確率に、$\dfrac{L(L+1)}{2}$ をかけたもの」と言うことができる

何らかの積の形になっているものを、選び方として捉えなおす方法を「積の和典型」と呼ぶが、これもその一種かな。

「直前の方向」と合わせて6通りの状態を管理することで、$6 \times 6$ 行列で遷移を表現できる。

          直前↓直前→
          0 1 2 0 1 2
直前↓ 0  p p p q        空欄は0
       1    p p   q
       2      p     q
直前→ 0  q q q p
       1    q q   p
       2      q     p

この遷移行列を $A$ として、$A^{N-1}$ の「直前→、0個選択」状態から、 「直前↓、2個選択」「直前→、2個選択」状態への和が、$L$ 成分の答えとなる。

同様のことを $R$ 成分でも行うことで、最終的な答えとなる。
今度は→移動から2個選ぶことになる。

Python3

import numpy as np


def mod_pow(x, a, MOD):
    ret = np.identity(x.shape[0], dtype=np.int64)
    cur = x
    while a > 0:
        if a & 1:
            ret = ret @ cur % MOD
        cur = cur @ cur % MOD
        a >>= 1
    return ret


t = int(input())
MOD = 998244353
INV100 = pow(100, -1, MOD)
for _ in range(t):
    n, p = map(int, input().split())
    p = p * INV100 % MOD
    q = (1 - p) % MOD
    dp = np.array([
        [p, p, p, q, 0, 0],
        [0, p, p, 0, q, 0],
        [0, 0, p, 0, 0, q],
        [q, q, q, p, 0, 0],
        [0, q, q, 0, p, 0],
        [0, 0, q, 0, 0, p],
    ], dtype=np.int64)

    res = mod_pow(dp, n - 1, MOD)
    ans = res[0, [2, 5]].sum()

    dp = np.array([
        [p, 0, 0, q, q, q],
        [0, p, 0, 0, q, q],
        [0, 0, p, 0, 0, q],
        [q, 0, 0, p, p, p],
        [0, q, 0, 0, p, p],
        [0, 0, q, 0, 0, p],
    ], dtype=np.int64)
    res = mod_pow(dp, n - 1, MOD)
    ans += res[0, [2, 5]].sum()
    ans %= MOD

    print(ans)