AtCoder Regular Contest 174 C,E 問題メモ

AtCoder Regular Contest 174

速解き回のようで速解き回でない、少し速解き回のARC

C - Catastrophic Roulette

• $1～N$ の目が等確率で出るルーレットがある
• 先手と後手がゲームをする
  • 交互にルーレットを回す
  • 出た目が初めて出た目なら、罰金なし
  • 出た目が過去に出ていたら、罰金1円支払う
  • 全ての目が出たら終了
• 先手、後手それぞれが支払う罰金の期待値を、$\mod{998244353}$ で求めよ
• $1 \le N \le 10^6$

まだ出てない目が $i$ 個で自分に手番が回ってきたときの、ここからの自分の罰金期待値 $P_i$、相手の罰金期待値 $Q_i$ とする。

最終局面は $P_0=0,Q_0=0$ である。

$i-1$ と $i$ の関係は以下のような式になる。

• 自分が回して出た目が、初めての目の場合（確率 $\dfrac{i}{N}$）
  • 自分は、相手に $i-1$ 個で渡すので、そこからの「相手の相手＝自分」の罰金期待値は $Q_{i-1}$
  • 相手は、$P_{i-1}$
• 既に出た目の場合（確率 $\dfrac{N-i}{N}$）
  • 自分は、罰金1円支払い、相手に $i$ 個で渡す。$Q_i+1$
  • 相手は、$P_i$

まとめると、

• $P_i=\dfrac{i}{N}Q_{i-1}+\dfrac{N-i}{N}(Q_i+1)$
• $Q_i=\dfrac{i}{N}P_{i-1}+\dfrac{N-i}{N}P_i$

ここから、$Q_i$ を消すと、$P_i$ を $P_{i-1},Q_{i-1}$ だけで計算できる式になる。それをもって $Q_i$ も計算できる。

$i$ が小さい方から求めていって、$P_N,Q_N$ が答え。

Python3

n = int(input())
MOD = 998244353

INV_N = pow(n, MOD - 2, MOD)
INV_N2 = INV_N * INV_N % MOD
p = 0
q = 0
for i in range(1, n + 1):
    np = i * INV_N * q % MOD
    np += (n - i) * i % MOD * INV_N2 * p % MOD
    np += (n - i) * INV_N % MOD
    np %= MOD
    np *= n * n * pow((2 * n * i - i * i) % MOD, MOD - 2, MOD)
    np %= MOD

    nq = i * INV_N * p % MOD + (n - i) * INV_N * np % MOD
    nq %= MOD

    p = np
    q = nq

print(p, q)

E - Existence Counting

• 「$1,...,N$ から $K$ 個取り出して並べる」ことで作られうる数列を「よい数列」とする
  • よい数列は ${}_NP_K$ 個ある
• よい数列の1つ $P=(P_1,...,P_K)$ が与えられる
• $t=1,2,...,N$ について、以下の条件を満たすよい数列の個数を $\mod{998244353}$ で求めよ
  • 整数 $t$ が含まれる
  • 辞書順で $P$ 以下である
• $1 \le K \le N \le 3 \times 10^5$

$P$ と先頭から比較してはじめて異なる値が出てくる位置を $d$ とする。
桁DPのように、$d$ ごとに分けて考えれば重複を除け、また $P$ より小さい条件を表しやすい。
$d$ には必ず $P_d$ より小さい値が置かれ、その後に置ける数は自由となる。

$t$ を固定すると、条件を満たす数列は以下の3つの場合がある。
$t$ が $P$ 中に出てくる位置を $u$ とする。$t$ が $P$ 中に出てこない場合、便宜的に $u=K$ とする。

• ① $t$ より左で $t$ より大きな値が出現する位置の1つを $d$ とし、$P_d$ を $t$ にする
  • $d$ より後に置く数は自由。その個数を数える
• ② $t$ または $t$ より左の位置を $d$ とし、$P_d$ を、$t$ 以外の $P_d$ より小さい値にする
  • $d$ より後に置く数は自由となるが、そのうちの1つが $t$ であるものの個数を数える
• ③ $u$ まで $P$ と一致する
  • 既に $t$ が含まれることは確定。$u+1$ 以降が自由である数列のうち、$P$ 以下の数列の個数を数える

N=18 K=13  t=10 の例
                      u
 i  0  1  2  3  4  5  6  7  8  9 10 11 12
 P  5 13 11  2 18  1 10 15 17  4 12  7  3

①  5 13 11  2 10  x  x  x  x  x  x  x  x
    5 10  x  x  x  x  x  x  x  x  x  x  x    など（xは任意）

②  5  9  x  x  x  x  x  x 10  x  x  x  x
    5 13 11  2 18  1  4  x  x  x 10  x  x    など    

③  5 13 11  2 18  1 10 15  6  x  x  x  x
    5 13 11  2 18  1 10 15 17  4 12  3  x    など

上記の例で、$d=4$ として18を10にした場合、その後に置ける箇所は $K-d-1=8$ 個、残る数字は $N-d-1=13$ 個ある。

よって、${}_{13}P_{8}$ 通りある。

このように、10より左で10より大きい箇所は、$d=1,2,4 \ (P_d=13,11,18)$ とあるので、合計すると以下の個数分ある。

• ${}_{16}P_{11}+{}_{15}P_{10}+{}_{13}P_{8}$

一般化すると、$t$ より左で、$t$ より大きい値が出現するindexの集合を $S_t$ として、①に該当する数列は以下の個数存在する。

• $\displaystyle \sum_{d \in S_t} {}_{N-d-1}P_{K-d-1}$

上記の例で、$d=1$ として、そこを $P_1=13$ より小さい何らかの値に変更したとする。

その後に置ける箇所は $K-d-1=11$ 個、残る数字は $N-d-1=16$ 個ある。
この中の1つが $t$ である数列の個数を数える。

$t$ を置く場所は11通り、他の箇所の決め方が ${}_{15}P_{10}$ 通り。
よって、$d=1$ の箇所に置く数を1つ決めた場合、以降に10が出てくる数列は、$11 \times {}_{15}P_{10}$ 通りある。

一般化すると、以下のようになる。（ただし、$K-d-2$ が負になる場合は $0$）

• $f(d)=(K-d-1) \times {}_{N-d-2}P_{K-d-2}$

$d$ の箇所に置ける数の個数分だけ、これが倍加される。

いま、$L_i$ を「$P_i$ 未満の数のうち、$i$ より左に出てこない数の個数」と定義する。
たとえば、$L_6$ なら、1,2,5 は $P_6=10$ より小さくて左に出てくるので、それ以外の 6 個が該当する。

「$d$ の箇所に置ける数」は、基本的に $L_d$ 個である。ただし、$t$ 自身であってはならないので、$t$ と $P_d$ の大小関係によって変化する。

• $P_d \le t$ の場合、$L_d$ 通りを $d$ に置ける
• $P_d \gt t$ の場合、$L_d-1$ 通りを $d$ に置ける

このように補正した値を $L'_d(t)$ とする。

②に該当する数列は、以下の個数分だけ存在する。

• $\displaystyle \sum_{d=0}^{u} L'_d(t) \times f(d)$

$t$ が $P$ 中に出てくる場合のみに考慮される。

要は、先頭のいくつかが固定されて、残りが自由である数列の中で、$P$ は何番目ですか、という値となる。

$i-1$ までが固定されて、$i$ 以降が自由であるよい数列の中で $P$ 以下のものの個数を $R_i$ とする。

i=10, Pi=12 の場合
                                 i
i  0  1  2  3  4  5  6  7  8  9 10 11 12    ここまでが固定されて、
P  5 13 11  2 18  1 10 15 17  4 12  7  3    残り3箇所が自由なよい数列の内、
   ~~~~~~~~~~~~~~~~~~~~~~~~~~~~           ←P以下のものの個数は？

$i$ に12未満の数を置いて $P$ 未満であることを確定させるか、または12のままにするか。
$i$ に置ける数の個数は、先述の $L_i$ と一致する。その後に置く数は自由となる。
12のままにした場合、その個数は $R_{i+1}$ と一致する。

まとめると、以下のようになる。

• $R_i = L_i \times {}_{N-i-1}P_{K-i-1} + R_{i+1}$

末尾から求めていける。

③に該当する数列は、$t$ が $P$ 中に出現する場合、$R_{u+1}$ 個存在することになる。（出現しない場合は0）

$t$ が小さい方から計算する。

①②は、$t$ によって加算される値が異なってくるが、

• 全て「自分より左の○○の合計」を求める
• 加算される値は①も②も「$P_i$ と $t$ の大小関係」に依存する

という性質があるため、Fenwick Tree で、$t$ が小さい方から答えを求めて $t$ の位置を更新していけば、管理できる。

t=1 のとき
i  0  1  2  3  4  5  6  7  8  9 10 11 12    ①: ①の場合の加算分
P  5 13 11  2 18  1 10 15 17  4 12  7  3    ②: ②の場合の加算分,係数 Li-1
  ① ① ① ① ①                            ❷: ②の場合の加算分,係数 Li
  ② ② ② ② ② ❷

t=2 のとき
i  0  1  2  3  4  5  6  7  8  9 10 11 12
P  5 13 11  2 18  1 10 15 17  4 12  7  3
  ① ① ①
  ② ② ② ❷

t=3 のとき
i  0  1  2  3  4  5  6  7  8  9 10 11 12
P  5 13 11  2 18  1 10 15 17  4 12  7  3
  ① ① ①    ①    ① ① ① ① ① ①
  ② ② ② ❷ ② ❷ ② ② ② ② ② ② ❷

t=4 のとき
i  0  1  2  3  4  5  6  7  8  9 10 11 12
P  5 13 11  2 18  1 10 15 17  4 12  7  3
  ① ① ①    ①    ① ① ①
  ② ② ② ❷ ② ❷ ② ② ② ❷

...

• 初期化
  • ①の個数
    • $d=0,...,k-1$ につき、$d$ に ${}_{N-d-1}P_{K-d-1}$ を加算
  • ②の個数
    • $d=0,...,k-1$ につき、$d$ に $(L_d-1) \times (K-d-1) \times {}_{N-d-2}P_{K-d-2}$ を加算

$t$ の答えを求める際は、$t$ の出現位置を $u$ として、まず自身の箇所を更新する。

• 更新
  • ①を削除
    • これ以降、$t$ の位置は加算されなくなる
    • $u$ から ${}_{N-u-1}P_{K-u-1}$ を減算
  • ②→❷に更新
    • $L_d$ に自身が含まれるため -1 していたものが、しなくてよくなる
    • $u$ に $(K-u-1) \times {}_{N-u-2}P_{K-u-2}$ を加算

（$t$ が $P$ 中に出てこない場合、便宜的に $u=K$ とするが、更新は行わないでよい）

この上で、FenwickTree 上の $0,1,...,u$ の累積和が①+②の場合、$R_{u+1}$ が③の場合に相当する。

$O(N \log{N})$ で全ての $t$ に対して求められる。

Python3

from operator import add


class FenwickTreeInjectable:
    def __init__(self, n, identity_factory, func):
        self.size = n
        self.tree = [identity_factory() for _ in range(n + 1)]
        self.func = func
        self.idf = identity_factory
        self.depth = n.bit_length()

    def add(self, i, x):
        i += 1
        tree = self.tree
        func = self.func
        while i <= self.size:
            tree[i] = func(tree[i], x)
            i += i & -i

    def sum(self, i):
        i += 1
        s = self.idf()
        tree = self.tree
        func = self.func
        while i > 0:
            s = func(s, tree[i])
            i -= i & -i
        return s

    def lower_bound(self, x, lt):
        """
        累積和がx以上になる最小のindexと、その直前までの累積和

        :param lt: lt(a, b) で a < b ならTrueを返す関数
        """
        total = self.idf()
        pos = 0
        tree = self.tree
        func = self.func
        for i in range(self.depth, -1, -1):
            k = pos + (1 << i)
            if k > self.size:
                continue
            new_total = func(total, tree[k])
            if lt(new_total, x):
                total = new_total
                pos += 1 << i
        return pos + 1, total

    def debug_print(self):
        prev = 0
        arr = []
        for i in range(self.size):
            curr = self.sum(i)
            arr.append(curr - prev)
            prev = curr
        print(arr)


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, k = map(int, input().split())
ppp = list(map(int, input().split()))
ppp = [p - 1 for p in ppp]
MOD = 998244353
d = n - k

facts, finvs = precompute_factorials(n, MOD)
npr = []
for i in range(k + 1):
    npr.append(facts[d + i] * finvs[d] % MOD)

pos = [-1] * n
for i, p in enumerate(ppp):
    pos[p] = i

fwt = FenwickTreeInjectable(k + 1, int, add)
i2self_lt = [-1] * k  # Pを先頭から見ていって、iまでに使われていない、Piより小さい値の個数
for p in range(n):
    i = pos[p]
    if i == -1:
        continue
    s = fwt.sum(i)
    i2self_lt[i] = p - s
    fwt.add(i, 1)

orders = [1] * (k + 1)  # i までprefixが同じ数列のうち、Pは何番目か
for i in range(k - 1, -1, -1):
    l = i2self_lt[i]
    orders[i] = l * npr[k - i - 1] + orders[i + 1]
    orders[i] %= MOD

medium = []  # i を Pi より小さい値に変えるとそれより後を自由に決められるが、その時に残り中に特定の1つの値が出てくるようなものの個数
for i in range(k):
    l = i2self_lt[i]
    if k - i - 2 >= 0:
        medium.append(npr[k - i - 2] * (k - i - 1) % MOD)
    else:
        medium.append(0)

fwt = FenwickTreeInjectable(k + 1, int, lambda a, b: (a + b) % MOD)
for i in range(k):
    fwt.add(i, npr[k - i - 1])
    if i2self_lt[i] >= 2:
        fwt.add(i, (i2self_lt[i] - 1) * medium[i] % MOD)

ans = [0] * n
for p in range(n):
    i = pos[p]
    if i == -1:
        tmp = fwt.sum(k)
        ans[p] = tmp
        continue

    fwt.add(i, -npr[k - i - 1])
    if i2self_lt[i] >= 1:
        fwt.add(i, medium[i])

    tmp = fwt.sum(i)
    tmp += orders[i + 1]
    ans[p] = tmp % MOD

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