AtCoder Japan Open -予選-（AtCoder Regular Contest 186）A,B,C,D問題メモ
- A - Underclued
  - 問題文
  - 制約
  - 解法
- B - Typical Permutation Descriptor
  - 問題文
  - 制約
  - 解法
- C - Ball and Box
  - 問題文
  - 制約
  - 解法
- D - Polish Mania
  - 問題文
  - 制約
  - 解法

AtCoder Japan Open -予選-（AtCoder Regular Contest 186）A,B,C,D問題メモ

AtCoder Japan Open -予選-（AtCoder Regular Contest 186）

A～Dが同じ配点で、どの問題を解くか、選球眼が試された。

A - Underclued

問題文

以下の問題を考えます。
- 長さ $N$ の2つの数列 $r=(r_1,...,r_N), c=(c_1,...,c_N)$ が与えられる。
- 各成分が $0,1$ である $N$ 次正方行列 $A$ のうち、以下の条件を満たす行列の集合を $S_{r,c}$ とする。
  - $i$ 行目の総和が $r_i$（$i=1,2,...,N$）
  - $j$ 列目の総和が $c_j$（$j=1,2,...,N$）
- $S_{r,c}$ 内の全ての行列で $A_{i,j}$ が同一（全て0または全て1）である場合、「$(r,c)$ に対して $(i,j)$ は固定されている」という。
- $f(r,c)$ を、「$(r,c)$ に対して固定されている要素の数」とする。

$N$ が与えられます。
以下の $Q$ 個のクエリに答えてください。
- $i$ 番目のクエリ：$f(r,c)=K_i$ となるような $(r,c)$ が存在するなら Yes、そうでないなら No を出力せよ。

制約

$2 \le N \le 30$
$1 \le Q \le N^2+1$
$0 \le K_i \le N^2$
$K_i \ne K_j (1 \le i \lt j \le Q)$
入力はすべて整数

解法

ある $r,c$ に対し、各行・各列の総和が $r,c$ と合致している行列の1つを $A$ とする。

$A_{i,j}$ を $0→1$ にしたら、辻褄を合わせるため、$i$ 行目・$j$ 列目のそれぞれ別の場所を、$1→0$ にする必要がある。

連鎖的に波及していき、どこかで閉じてループにすれば、それらを一斉に反転させることで総和を変えずに値が入れ替わる。
このような、「0-1を交互に拾い、拾ったら必ず90度曲がるループ」の一部になり得ない要素が、固定された要素となる。

0-----1         1-----0
|     |         |     |
|     0--1  →  |     1--0
1--0     |      0--1     |
   |     |         |     |
   1-----0         0-----1

ところで、以下のように長方形の4頂点中、3頂点がループに含まれるものの残りの1頂点は、最初が0であろうと1であろうと、そこで新たなループが作れてしまう。

0--!--1       1--!--0        ! の部分は、1なら左、0なら右の場合に
|  :  |       |  :   |       新たなループに含まれる。
|  :  0--1    |  :  1--0
1--0     |    0--1     |
   |     |       |     |
   1-----0       0-----1

これを連鎖的に適用すると、ループ上の点と同じ行×同じ列に含まれる頂点は全て、ループに含まれうる（＝固定されていない）要素となる。

o--o--o--o
|  |  |  |
o--o--o--o
o--o--o--o
|  |  |  |
o--o--o--o

行ごと・列ごとの入れ替えは（それに伴い $r,c$ も並び替えれば）総和に影響しないので、固定されてないのを左上に集める。

????ffffff
????ffffff    ?: 自由（固定されていない）
????ffffff    f: 固定されている
ffffxxxxxx    x: 不明
ffffxxxxxx
ffffxxxxxx

4領域に分かれる。
ここで、自由領域（?がある箇所）の幅・高さは、ループを構成するのでそれぞれ $2$ 以上である必要がある。

右下には不明領域が残るが、「ここも全て固定されていなくてはならない」のか、「固定されていなくてもいい（既存の固定領域に影響はない）」のか、どっちだろうか？

ループの最小要素である、0-1を交互に拾う4点の「長方形ループ」を作れるかどうかで考える。
固定領域の頂点が含まれるようなループが存在してはいけない。

影響があるとすると、4領域からそれぞれ1頂点ずつ選んで長方形を作った場合となる。
これがループとならないようにするには、固定領域の一方を0、一方を1で埋めればよい。

????111111    ? を0/1のいずれにしても、
???---1111    0-1を交互に巡るループにはなり得ない。
??|?11|111
00|0xx|xxx    また、? や x の中で複数点を通過するようなループでも、
000---?xxx    0-1を交互に90度曲がりながら拾う、というルールのため、
0000xxxxxx    固定領域の2箇所はいずれにしても同じ要素が来ないと成り立たない。

これで、右下の領域には何を置いても既存の固定領域には影響がないことが保証された。

右下の“x”領域の中身を、スケールの小さな問題に再帰的に分割できる。

再帰的に考えた結果、「固定されていない要素数」としてあり得る数は、以下のような値に限られることが分かる。

左上から、行・列を共有しないように、幅・高さが2以上の長方形を置いていく。
- 必ずしも、右下まで長方形で埋め切る必要はない点に注意。
長方形の面積の合計が「固定されていない要素数」として取り得る値
$N^2$ から引いた値が、「固定された要素数」として取り得る値

■固定されていない領域の決め方の一例

????1111111
????1111111    ? :固定されていない領域
0000??11111
0000??11111
0000??11111
000000???11
000000???11
00000000011

$\mathrm{DP}[n,m]:=$ $n$ 行 $m$ 列の時に、固定されていない要素数として取り得る値の集合

として $n,m$ の小さい方から求めていく（またはメモ化再帰で実装する）ことで、DPできる。
状態 $O(N^2)$、遷移 $O(N^2)$、集合をマージするのに $O(N^2)$ かかり、全体 $O(N^6)$ 程度で可能となる。（定数倍で数分の1になると見込める）

Python3

def partition(n, m, lo=2, memo={}):
    if n <= 1 or m <= 1 or n < lo:
        return {0}
    key = (n, m, lo)
    if key in memo:
        return memo[key]
    result = {0}
    for i in range(lo, n + 1):
        for j in range(2, m + 1):
            res = partition(n - i, m - j, i)
            ij = i * j
            result.update(ij + p for p in res)
    memo[key] = result
    return result


n, q = map(int, input().split())
pt = partition(n, n)
n2 = n * n
for _ in range(q):
    k = int(input())
    if n2 - k in pt:
        print('Yes')
    else:
        print('No')

B - Typical Permutation Descriptor

問題文

$1,...,N$ の順列 $P=(P_1,\dots,P_N)$ に対して、以下で定義される整数列 $B$ を $f(P)$ とします。
- $B_i$ は「$P$ 上で $i$ より左で、$P_i$ より小さい数が現れる直近のindex」を示す。
- 存在しない場合は $B_i=0$ とする。
- より厳密には、$i=1,\dots,N$ について、
  - $0\le B_i \lt i$
  - $B_i \lt j \lt i$ であるすべての整数 $j$ について、$P_j \gt P_i$
  - $B_i \gt 0$ ならば $P_{B_i} \lt P_i$

長さ $N$ の整数列 $A$ が与えられるので、$f(P)=A$ となるような順列 $P$ の数を $998244353$ で割ったあまりを求めてください。
ただし、この問題の入力で与えられる $A=(A_1,\dots,A_N)$ について、条件を満たす順列が存在することが保証されます。

制約

$1\le N\le 3\times 10^5$
$0\le A_i \lt i$
$A_1,\dots,A_N$ について、問題文中の条件を満たすような順列が存在する
入力はすべて整数

解法

問題文のとおりに大小関係を有向グラフ（大→小）で表してみる。

グラフの辺を引く基準
- $A_i$ ← $i$
- $i$ より右で、$i$ を追い越して $i$ より左に戻るような $j$（$A_j \lt i \lt j$）は、$i$ → $j$

として構築していける。
$P_0=0$ 固定である頂点⓪を仮想的に追加して、

i   1  2  3  4  5  6  7  8
A   0  1  1  3  0  5  5  7

     ,--------------ｖ
⓪←①←②→③←④  ⑤←⑥→⑦←⑧
 ↑ ＾------'       |^------'
  `-----------------'             ※自明な辺は一部略

答えは、このグラフのトポロジカルソートの数え上げとなる。
だが、通常のトポロジカルソートの数え上げは $O(N2^N)$ とかになって無理。何らかの計算しやすい性質があるはず。

このグラフをよく見ると、不要な辺を除くと、必ず有向木になる。

⓪←①の辺は、⓪←⑤←①があるので不要、など

⓪←⑤←①←③←②
    ↑      ＾--④
     `--⑦←⑥
        ＾--⑧

略証

各頂点につき、「$i←k$ かつ $j←k$ なら、$i$ と $j$ の間にも必ず大小関係がある」ことを言えればよい。
それが言えると、$k$ から出る辺のいずれかは不要となり、各頂点から出る辺は必ず1本となる＝有向木となることが示せる。

「グラフの辺を引く基準」のいずれで $i←k,j←k$ の辺が引かれたかを場合分けすると、

① Ak = i であるケース（Ak=j も対称的に考えると同様）
   j は k を追い越して戻ることになる。Aj と i の位置関係で場合分けする。
   
  Aj=i          Aj<i             Aj>i
                                 v-------,
  i ← k   j    o  i ← k   j    i   o   k   j
  ^--------'    ^-----------'        ^-------'
                                 ×制約に矛盾
  左2つのいずれかとなる。
  いずれの場合も、iとj の間にも大小関係が定義される。

② それ以外のケース
   i,j ともに k を追い越して戻ることになる。

  o o  k  i  j    o  k  i  j
  ^ ^-----'  |    ^-----'  |
  `----------'    `--------'
  いずれにしても iとj の間に大小関係は定義される。

木なら、木DPでトポロジカルソートの数え上げができる。

$ans=1$ からはじめて、複数の子を持つ頂点で子をマージするときに、各子の部分木のサイズを $s_1,s_2$ として、${}_{s_1+s_2}\mathrm{C}_{s_1}$ を $ans$ にかけていけばよい。

木の構築は、

$A_i=i$ である $i$ を $A_i$ ごとに列挙する。$I(A_i)=(i_1,i_2,...,i_k)$
$A_i$ が小さい方から順に、$A_i←i_k←i_{k-1}←...←i_2←i_1$ という辺を追加する。

としていくとできる。

Python3

from collections import defaultdict


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


def solve(n, aaa):
    MOD = 998244353
    facts, finvs = precompute_factorials(n + 1, MOD)

    pos = defaultdict(list)
    for i, a in enumerate(aaa, start=1):
        pos[a].append(i)

    q = [0]
    links = [[] for _ in range(n + 1)]
    while q:
        u = q.pop()
        for v in reversed(pos[u]):
            links[u].append(v)
            q.append(v)
            u = v

    dp = [-1] * (n + 1)
    q = [0]
    ans = 1
    while q:
        u = q[-1]
        if dp[u] == -1:
            dp[u] = -2
            for v in links[u]:
                q.append(v)
        else:
            q.pop()
            cnt = 0
            tmp = 1
            for v in links[u]:
                tmp *= facts[cnt + dp[v]] * finvs[cnt] * finvs[dp[v]] % MOD
                tmp %= MOD
                cnt += dp[v]
            ans *= tmp
            ans %= MOD
            dp[u] = cnt + 1

    return ans


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

C - Ball and Box

問題文

球橋さんと箱木さんはボールと箱を使ったゲームをします。
最初、球橋さんは $M$ 種類のボールをそれぞれ $10^{100}$ 個ずつ持っていて、箱木さんは $10^{100}$ 円持っています。
また、$N$ 個の箱があり、$i$ 番目の箱の容量は $V_i$ で、値段は $P_i$ 円です。
- ゲーム中、箱木さんはいつでも好きな箱を買うことができます。
このゲームでは、ゲームが終わるまで以下の操作を繰り返します。
1. 球橋さんはボールを $1$ つ選び、箱木さんに渡す。
2. 箱木さんは渡されたボールを受け取るか、受け取らずにゲームを終えるかを選ぶ。
3. ボールを受け取った場合、箱木さんは購入済みの箱を $1$ つ選び、受け取ったボールを入れる。
4. ボールを入れた箱が以下の条件を満たしている場合、箱木さんは $1$ 円をもらう。そうでない場合、ゲームを終える。
  - 箱の中のボールの個数は、その箱の容量以下である。
  - 箱の中のボールの種類は、すべて同じである。

球橋さんは、ゲーム終了時の箱木さんの所持金がなるべく少なくなるための最適な行動をし、反対に、箱木さんはなるべく多くなるための最適な行動をします。
ゲームを通して、箱木さんの所持金はいくら増えますか？
ただし、両者ともにすべての情報が公開されているとします。特に、球橋さんは、それぞれの箱について、容量と値段、どの種類のボールがいくつ入っているかを見ることができます。
また、箱木さんの初期の所持金は十分多く、お金が足りなくて箱が買えなくなることはないことに注意してください。
$1$ つの入力ファイルにつき、$T$ 個のテストケースを解いてください。

制約

$1\le T,N,M\le 3\times 10^5$
$1\le V_i,P_i \le 10^9$
$T$ 個のテストケースに対する $N$ の総和は $3\times 10^5$ 以下
入力はすべて整数

解法

2人の行動について「短期的に最適にするにはこういう戦略がよさそうだな」というのは実験してみれば何となくわかり、実際その考えを進めていくとACはできるが、「本当に大局的に見てそうなの？」の証明はなかなか難しい。

公式Editorialにならい「種類 $a$ のボール」を「色 $a$ のボール」と呼ぶことにする。
また、「既に色 $a$ のボールが入っている、容量が余っている箱」を「色 $a$ の箱」とする。

最初に買う箱を決める場合

まず、仮に箱木さんが最初に買う箱 $B=(B_1,B_2,...,B_k)$ を決め、後は一切買わないとする。

購入費用は確定済みなので、気にするべきはボール数のみになる。
球橋さんはなるべく少ないボール数で、箱木さんはなるべく多いボール数でゲームを終了させたい。
ゲームは「ある色のボールにつき、もう入れられる箱が無い」状態で、球橋さんがその色のボールを渡すと終了する。

球橋さんは、短期的には以下の行動をするのが最適であるように見える。これを「戦略 Tama」とする。

戦略 Tama
- 色 $a$ の箱が無ければ、色 $a$ のボールを渡す。
  - 箱1個の色を確定させ、今後の箱木さんの自由度を下げられる。
- 全ての色 $a$ について箱が存在するなら、箱の残容量が最も少ない色のボールを渡す。
  - なるべくはやく「ある色のボールにつき、もう入れられる箱が無い」にする。
  - 色 $a$ の箱が複数ある場合、全ての箱の残容量の合計を「色 $a$ の残容量」とする。

箱木さんは以下の戦略をとるのが良さそうに見える。これを「戦略Hako」とする。

戦略 Hako
- 既に色 $a$ の箱がある状態で色 $a$ のボールが渡されたら、既存の箱に重ねて入れる。
- 箱を選択する順番はどうでもよい。

両者が最適に行動すれば、最終的には以下の状態（★とする）になる。

最終状態（★）
- $V_{B_1},V_{B_2},...,V_{B_k}$ のうち容量が大きい方から $M-1$ 個の箱には、それぞれ異なる色のボールが1個ずつ入っている。
- それ以外の箱には、容量分のボールが入っている。

ただし、購入した箱数 $k \le M-1$ の場合は、全ての箱に1個のボールが入った状態で終了する。
この場合は、$P_i \ge 1$ よりどうしても利益を上げることができないので、何もしない方がましとなる。
以下、$k \gt M-1$ とする。

最終状態が（★）になる略証

（★）における箱木さんの儲けを $S$ とする。

①球橋さんが戦略Tamaを取る前提のもと、箱木さんは「戦略 Hako で必ず（★）を達成できること」「他に儲けが $S$ を超えるような最終状態にはできないこと」を示す。
②箱木さんが戦略Hakoを取る前提のもと、球橋さんは「戦略 Tama で必ず（★）を達成できること」「他に儲けを $S$ より少なくするような最終状態にはできないこと」を示す。

①と②をともに示すことで証明できる。2人の最適な戦略も確定する。

①球橋さんが戦略Tamaをとる前提で、箱木さんは（★）を達成できるか？よりよい最終状態にはできるか？

色 $a=1,2,...,M$ のボールを入れた箱の残容量を $U_a$ とする。そのような箱が無いときは $U_a=0$ とする。
未使用の箱の個数を $r$ とする。最初、$r=k,U_a=0 (a=1,2,...,M)$ である。
$r=0$ かつどれか1つの $a$ について $U_a=0$ にすればゲーム終了となる。

箱購入時点で全体の容量 $\sum V_i$ は決まっているので、最終的に使われずに余る $\sum U_a$ を少なくすることが、箱木さんの目標となる。

戦略Tamaに従うと、最初に色 $1～M$ のボールが1つずつ渡される。これらは箱木さんによりそれぞれ別の箱に入る。

（※）次に $U_a$ が最小の色の1つを $b$ として、ボール $b$ が渡される。
箱木さんはこれを「既存の箱に重ねて入れる（戦略Hako）」か「新たな箱に入れる」かの選択がある。

「既存の箱に重ねて入れる」場合、$U_b$ は1減る。元から $U_b$ が最小だったのにさらに1減ったので、続けて $b$ が渡され続ける。
やがて $U_b=0$ となり、さらに $b$ が渡される。
未使用の箱が残っていれば、その箱の1つを選んで入れる。選んだ箱を $i$ として、$U_b←V_i-1$ になる。
（※）に戻り、再度、現時点で $U_a$ 最小の色が特定され、未使用の箱が無くなるまで繰り返される。

戦略Hakoを採り続けると、最終状態は（★）となる。

「既に同じ色の箱があるが、新たな箱に入れた」場合を考える。
新たな箱を $j$ とすると、残容量は $U_b ← U_b + V_j-1$ となる。
ここで、（★）において「$U_a \gt 0$ のうち、最小のもの」を $U_c$ とする。
$U_b \gt U_c$ となった場合、戦略Tamaでは $U_a$ が大きいものを残そうとするので、最終状態に残る $M-1$ 個の中から、$U_c$ の代わりに $U_b$ が入ることとなる。

ざっくり言うと、最終的に残る $U_a$ の入れ替わりは $\sum U_a$ が増える方向でしか行われないので、（★）以外の最終状態があったとしても、儲けは $S$ より少ないといえる。

よって、球橋さんが戦略 Tama を取ってきた場合、箱木さんは戦略 Hako で（★）に持っていくのが最適となる。

②箱木さんが戦略Hakoを取る前提のもと、球橋さんは $S$ より儲けが少ない状態にできるか？

ゲーム終了の条件は、「どれか1つのボールにつき、もう入れられる箱が無い」状態である。
言い換えると、他の $M-1$ 色は容量が余っている箱があってもよい。

各 $M-1$ 色につき容量が余っている箱があり、「そのような箱が各色1つのみ」である中で、（★）は最も儲けが少ない状態である。
（$V_i$ が多い方から容量を余らせているし、それぞれに最低限の1個しか入れていないので）

また、箱木さんが Hako を取る以上、「ある色について、容量の余っている箱が2つ以上ある」ような状態にはならない。

よって、球橋さんは $S$ より儲けが少なくはできない。

箱の購入が最初でなくてもいい場合

最初にのみ箱を買う方針で、全ての「最初に買う箱の組」を固定した中で最大の儲けを $S$ とする。

「箱木さんが途中でも箱を買ってよい」とした場合に、儲けを $S$ より大きくできるか？

これはできない。

略証

よって、最初に買う箱を全て決めてしまう場合を考えてよく、$S$ を求めればよいことになる。

答えを求める

当然、全ての箱の買い方の全探索はできない。

（★）において容量上位 $M-1$ 個の箱以外は容量をフルで使えるので、$V_i \gt P_i$ なら必ず買って良い。

$V_i$ を大きい順にソートし、「①$1～i$ から $M-1$ 個の箱を、ボールが1個だけ入る箱として選ぶ」「②$i+1$ 以降で $V_i \gt P_i$ の箱は全て選ぶ」として、境界 $i$ を全探索すれば、$O(N)$ 通りの探索で済む。

①で、$M-1$ 個の箱は容量はもはや関係なくなるので、価格を基準に決めてよい。
ヒープキューなどで管理して、安い方から $M-1$ 個を買えばよい。

②では、末尾からの累積和を前計算しておけばよい。

Python3

from heapq import heapify, heappushpop

t = int(input())
buf = []
for _ in range(t):
    n, m = map(int, input().split())
    boxes = [tuple(map(int, input().split())) for _ in range(n)]
    boxes = [box for box in boxes if box[0] > box[1]]

    if m == 1:
        ans = sum(v - p for v, p in boxes)
        buf.append(ans)
        continue

    if len(boxes) <= m - 1:
        buf.append(0)
        continue

    boxes.sort(reverse=True)
    # v-p の末尾からの累積和
    all_get = [0]
    for v, p in reversed(boxes):
        all_get.append(all_get[-1] + v - p)
    all_get.reverse()

    sum_wasted_p = 0
    wasted = []
    for i in range(m - 1):
        v, p = boxes[i]
        sum_wasted_p += p
        wasted.append(-p)
    heapify(wasted)

    ans = m - 1 - sum_wasted_p + all_get[m - 1]
    for i in range(m - 1, len(boxes)):
        v, p = boxes[i]
        p2 = -heappushpop(wasted, -p)
        sum_wasted_p += p - p2
        ans = max(ans, m - 1 - sum_wasted_p + all_get[i + 1])
    ans = max(ans, 0)
    buf.append(ans)

print('\n'.join(map(str, buf)))

D - Polish Mania

問題文

空でない非負整数列 $(V_1, V_2, \dots, V_M)$ が Polish であることを、次のように再帰的に定義します。
- $V_1$ 個の Polish 数列 $W_1, W_2, \dots, W_{V_1}$ が存在して、数列 $(V_1), W_1, W_2, \dots, W_{V_1}$ をこの順に連結したものが数列 $(V_1, V_2, \dots, V_M)$ と一致する
特に、数列 $(0)$ は Polish です。
長さ $N$ の非負整数列 $(A_1, A_2, \dots, A_N)$ が与えられます。辞書順で $(A_1, A_2, \dots, A_N)$ 以下である、長さ $N$ の Polish 数列の数を $998244353$ で割ったあまりを求めてください。

制約

$1\leq N \leq 3\times 10^5$
$0\leq A_i \lt N$
入力はすべて整数

解法

Polish数列は、木に対応づけられる。
（答えを求める上でこの言い換えは必ずしも必要ないが、その後の性質を発見しやすくなると思う）

P  3 2 0 2 1 0 0 1 3 0 0 0 0

      ③       各頂点の数字は Pi に対応し、かつ子の数と一致する。
    ／｜＼     
   ② ① ⓪    子が複数ある場合は左から辿るというルールで、
 ／｜  ＼      オイラーツアーを行きがけ順に記録したものが P と一致する。
⓪ ②    ③
   /\   /｜\   木とPolish数列は、いずれの方向にも一意に変換でき、1対1対応する。
  ①⓪ ⓪⓪⓪
  ｜
  ⓪

ある頂点の「部分木内に書かれた数の総和」は、「部分木内の自身を含まない頂点数」と等しくなる。
（葉から根に向けて、子の数を足していってるのでまぁ、当然と言えば当然）

特に、全体の総和は $N-1$ となる。

では、全体の総和が $N-1$ ならどんな数列でもPolishかというと、当然そうはならない。
たとえば、末尾は必ず 0 でなければならない。オイラーツアーの行きがけ順の最後は必ず葉が来て、葉は必ず⓪になるので。

条件を考えると、$P$ を逆にした数列 $Rev(P)$ に対して、以下が成り立てばよいことがわかる。

各 $i=1,2,...,N$ に対して、$\sum_{k=1}^{i}Rev(P)_k \lt i$

つまり、$i$ までの累積和が $i$ 未満であればよい。数列がこの条件を満たしている場合、以下のように木を必ず構築できる。

空のスタックを用意する。
$Rev(P)$ の先頭から、要素を見ていく。
要素が0の場合、スタックに1頂点からなる木 $⓪$ を積む。
要素 $p \ge 1$ の場合、スタックから $p$ 個の木を取り出し、自身にそれらを子として繋げて1つの木とし、スタックに積む。

なので（辞書順のことはいったん置いておいて、方針としては）このような条件を満たす数列 $Rev(P)$ の個数を求めたい。

これは実は、カタラン数で表せる。
グリッド上を斜めの線を越えないように移動する方法と、累積和数列が、1対1対応する。

カタラン数

     4       ／
     3     ／    →→ ...
     2   ／→→→↑
     1 ／  ↑
     0 →→↑

Rev(P)  0 0 2 0 0 1 0 ...

制約がなければ、長さ $N$ のPolishな数列はカタラン数 $C_{N-1}$ で数えられることがわかった。

では、「与えられた $A$ より辞書順で小さい」を考慮するとどうなるか。

$P$ の先頭からの制約が生じるので、カタラン数のグリッドを右側から見ていく。
経路も、$t$ から遡って考えていくこととする。

     j
    10                   ／ t
     9                 ／
     8               ／
     7             ／       -
     6           ／   - - - x
     5         ／     x x x x
     4       ／       x x x x
     3     ／       - x x x x
     2   ／         x x x x x
     1 ／         - x x x x x
     0 s        - x x x x x x
i       1 2 3 4 5 6 7 8 91011
Rev(A)  5 0 4 0 1 2 3 0 0 1 3

$j=N-1=10$ からスタートし、末尾から $Rev(A)$ を累積的に減算していったものが “-“、それより下が “x” で表されている。

「”-” まで下がってから1つ左の列に移る」ことを繰り返している間は、次の列でも、“-” より下に移動してはいけない。
「“-” より上で1つ左の列に1度でも移った」場合、それ以降は “-”,“x” の制約はなくなり、“／”さえ超えなければよい。

たとえば $Rev(A)_{11}=3$ より、$i=11$ の列では $j=10$ から $3$ つを超えて下がってはいけない。
$i=11$ で $j=7$ まで下がってから $i=10$ の列に移るような経路を取った場合、$i=10$ でも $Rev(A)_{10}=1$ つを超えて下がってはいけない。

一方、$i=11$ で2つ以下しか下がらない内に $i=10$ の列に移った場合は、辞書順で $A$ 未満であることが確定し、残りは自由となる。
以降は“／”の制約のみ満たせばよいことになる。

どこで $i=10$ に移ったかで場合分けした図で表すと、

◎の場合は“x”の制約がなくなる。
- 制約が“／”のみになるため、まとめて計算できる。経路数を求めて答えに加算する。
- この経路数は、／を無視して「s から◎への経路数 - s から●への経路数」で求められる。
  - グリッドの最短経路の数え上げまとめ - かんプリンの学習記録
- さらに、同じ列の◎,●への経路はまとめて計算できる。
  - Wolfram Alpha
○の場合は、1列短くなって同様の議論を繰り返す。

    10               ●●／ t
     9                 ／◎←
     8               ／  ◎←
     7             ／    ○←
     6           ／         x
     5         ／     x x x x
     4       ／       x x x x
     3     ／         x x x x
     2   ／         x x x x x
     1 ／           x x x x x
     0 s          x x x x x x
i       1 2 3 4 5 6 7 8 91011
Rev(A)  5 0 4 0 1 2 3 0 0 1 3

“x” が “／” と同じ高さまで迫り、置ける数がなくなったら、それ以上、prefixが $A$ と同じで $A$ 以下のPolishな数列は作れないので、打ち切って良い。

“x” の高さが負まで下がったら、あとはPolishであれば必ず $A$ 以下になるので、残りの（“x”の制約のない）経路数を答えに加算した後、打ち切って良い。

$A$ 自体がPolishである場合、最後に1を足すのを忘れないように注意。

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


def solve(n, aaa):
    MOD = 998244353
    facts, finvs = precompute_factorials(n * 2, MOD)

    def partial_catalan(w, hi, lo):
        if w < hi:
            return 0
        res1 = (hi + 1) * facts[w + hi + 1] * finvs[w] * finvs[hi + 1] % MOD
        res1 -= lo * facts[w + lo] * finvs[w] * finvs[lo] % MOD
        res1 *= facts[w] * finvs[w + 1] % MOD
        if lo > 0:
            hi -= 1
            lo -= 1
            w += 1
            res2 = (hi + 1) * facts[w + hi + 1] * finvs[w] * finvs[hi + 1] % MOD
            res2 -= lo * facts[w + lo] * finvs[w] * finvs[lo] % MOD
            res2 *= facts[w] * finvs[w + 1] % MOD
            res1 -= res2
        res1 %= MOD
        return res1

    ans = 0
    lo = n - 1
    for i, a in enumerate(aaa[:-1]):
        cur_lo = lo - a
        w = n - i - 2
        if cur_lo > w:
            break
        if cur_lo < 0:
            ans += partial_catalan(w + 1, lo, lo)
            ans %= MOD
            break
        ans += partial_catalan(w, min(lo, w), cur_lo + 1)
        ans %= MOD
        lo = cur_lo
    else:
        ans += 1
        ans %= MOD
    return ans


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

目次

AtCoder Japan Open -予選-（AtCoder Regular Contest 186）A,B,C,D問題メモ

A - Underclued

問題文

制約

解法

B - Typical Permutation Descriptor

問題文

制約

解法

C - Ball and Box

問題文

制約

解法

最初に買う箱を決める場合

箱の購入が最初でなくてもいい場合

答えを求める

D - Polish Mania

問題文

制約

解法