AtCoder Beginner Contest 226 F,G,H 問題メモ

AtCoder Beginner Contest 226

PythonistaはさぞかしA問題でWAした人が多いことでしょう。
偶数丸め、忘れがち。

F - Score of Permutations

• $N$ 人の人と $N$ 個のボールがあり、はじめ、人 $i$ はボール $i$ を持っている
• 長さ $N$ の順列 $P=(p_1,...,p_N)$ に対し、以下でスコア $S(P)$ を定義する
• スコア
  • 号令を1回かけると、人 $i$ は自分が持っているボールを人 $p_i$ に同時に渡す
  • 全ての人が自分の番号のボールを持つ（初期状態に戻る）までの回数をスコアとする
• $P$ としてあり得るものは $N!$ 通りあるが、全ての $P$ について「$S(P)$ を $K$ 乗した値」の総和を求めよ
• $2 \le N \le 50$

$K$ 乗させることの意味をいまいちわかってない……
と思ったら、入力が $N$ だけだと手元で時間かけて計算して埋め込めるからか、なるほど。

このように順列を $i→p_i$ に置換していく問題は過去に何度も出ているが、 $i→p_i→p_{p_i}→...→i$ と追っていき $i$ に戻るまでを1つのサイクルとすると、 順列はいくつかのサイクルに分割できる（長さ1もサイクルとする）。

1つのサイクルだけを見ると、それが何回ごとに元に戻るかの周期は、サイクルの長さと一致する。

よって $P$ を固定した時、$P$ が長さ $L_1,L_2,...$ のサイクルから成っていたとして、 スコアは $L_i$ 全体の最小公倍数となる。

スコアが $x$ となる $P$ の個数を $x$ ごとに数えられたら、答えは求められる。

• $DP[i][j]=i$ 人のサイクルを決めて、その際の暫定スコアが $j$ である場合の数
  • 暫定スコアとは、既に決めたサイクルのLCMとする。

$DP[i][j]$ から新たに長さ $l$ のサイクルを追加すると、$DP[i+l][LCM(j,l)]$ に遷移することになる。

重複して数えないように注意が必要。たとえば

• 長さ3のサイクルを1つ決める（$DP[3][3]$）→ 長さ2のサイクルを追加して $DP[5][6]$ に遷移
• 長さ2のサイクルを1つ決める（$DP[2][2]$）→ 長さ3のサイクルを追加して $DP[5][6]$ に遷移

前者で $(1,3,4),(2,5)$、後者で $(2,5),(1,3,4)$ をそれぞれ別に数えてしまうと、重複してしまう。

このような場合、 「新しく追加するサイクルには、残っている中で番号が最小のものは必ず入れる」 と考えると重複を除いて数えられる。

よって、$DP[i][j]$ から長さ $l$ のサイクルを追加する遷移は以下のようになる。

• まだ決まってない残り人数は $N-i$、これを $r$ とする
• 1人は番号最小の人固定で、残り $r-1$ 人から $l-1$ 人を選ぶ
  • サイクルの並べ方も考慮すると、${}_{r-1}P_{l-1}$ の作り方がある
• $DP[i+l][LCM(j,l)] += DP[i][j] \times \dfrac{(r-1)!}{(r-l)!}$

$j$ の取り得る値の種類数（総和が $N$ 以下の複数の値のLCMとなりうる数）を $T$ として、$O(N^2T)$ で全てを計算できる。

$T$ のオーダー計算がちょっと難しいが、 $N=50$ で試しに実行しても1056程度であることが確認でき、十分高速であることがわかる。

全て計算できたら、各 $j$ につき、$DP[N][j] \times j^K$ の総和が答え。

Python3

from collections import defaultdict
from math import gcd


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())
MOD = 998244353
facts, finvs = precompute_factorials(n, MOD)

dp = [defaultdict(int) for _ in range(n + 1)]
dp[0][1] = 1
for i in range(n):
    rem = n - i
    for j, c in dp[i].items():
        for l in range(1, rem + 1):
            lcm = j * l // gcd(j, l)
            dp[i + l][lcm] += c * facts[rem - 1] * finvs[rem - l]
            dp[i + l][lcm] %= MOD

ans = 0
for j, c in dp[n].items():
    ans += c * pow(j, k, MOD)
    ans %= MOD

print(ans)

G - The baggage

Gだからそこまでたどり着かなかったり変に難しく考えた人が多いだけで、DとかEに置かれてたらそこそこ（未証明を含め）解かれてそうな気がする？

• 重さが $1,2,3,4,5$ の荷物がそれぞれ $A_1,A_2,A_3,A_4,A_5$ 個ある
• 体力が $1,2,3,4,5$ の人がそれぞれ $B_1,B_2,B_3,B_4,B_5$ 人いる
• 体力が $h$ の人は、重さの合計が $h$ 以下の複数の荷物を持つことが出来る
• 全ての荷物を誰かが持っている状態にすることが可能か、判定せよ
• 1つの入力につき $T$ 個のテストケースが与えられる
• $1 \le T \le 5 \times 10^4$
• $0 \le A_i,B_i \le 10^{16}$

重さの合計を $W_A,W_B$ として、$W_A \gt W_B$ は問題外として、 $W_A \le W_B$ でも不可能なケースがある。

重さ 2 の荷物が 5 個
体力 3 の 人 が 4 人  答えはNo

重さ 4 の荷物が 1 個
体力 3 の 人 が 9 個  答えはNo

貪欲に決められないか考える。

なるべく持てる人が限られる重い荷物から決めていきたい。
また、人は、体力の端数が出ないならばなるべく重い荷物を優先的に持った方がよい。

重さ $i$ の荷物はまず体力 $i$ の人に持たせて損しない。
「荷物 $i$ を持っていない体力 $i$ の人(A)がいるのに、$i$ 以外の体力で荷物 $i$ を持っている人(B)がいる」状態で成立している場合、 (A)の持つ全荷物を(B)の荷物 $i$ と交換しても成立は保たれる。

従って、$1～5$ について荷物と人のどちらか一方は相殺して0にでき、どちらかが余っている状態になる。

ここからは重い方の荷物の割り当てから考えていく。どこかのタイミングで、その荷物を持てる人がいなくなればアウト。

重さ5の荷物が余っていたら不可能。

重さ4の荷物が余っていたら、体力5の人がそれ以上余っていれば持たせられる。
荷物を持った人は、以降体力1として扱う。

ここまでは、そうするしか方法が無い。

重さ3の荷物が余っていたら、体力5か4の人に持たせられ、選択の余地が生まれる。
ここで、どちらに持たせるべきか？ 貪欲に判定可能か？

• 体力5に優先的に持たせる → 体力 2 4 が多く残る
• 体力4に優先的に持たせる → 体力 1 5 が多く残る

残る荷物の重さは 1 と 2 である。

1は最小単位なので、どのように体力を残しても総和が足りていれば持たせられる。
総和が足りていてもアウトになるのは、重さ2の荷物が残るのに、(残り)体力1の人しかいなくなった状態である。

奇数である1,5を多く残した場合、その分だけ端数（残り体力が1となる人）が発生しうる。 しかし、偶数の2,4を多く残せば、端数を最小限に抑えられる。

重さ3の荷物をどちらに優先的に持たせても残る体力の総和は変わらないので、体力 $2,4$ が多く残るようにした方がよいことがわかる。

実装上は、まずは同じ重さでマッチングした後、 「重さの大きい方の荷物から順に」「体力の大きい方の人から順に」割り当てていけばよい。

これは全体が $1～5$ だから成立するのであって、$6$ 以上では成立しない点に注意。
（その場合はフローで解ける？）

Python3

t = int(input())
for _ in range(t):
    aaa = [0] + list(map(int, input().split()))
    bbb = [0] + list(map(int, input().split()))

    ddd = [a - b for a, b in zip(aaa, bbb)]

    for i in range(5, 0, -1):
        if ddd[i] <= 0:
            continue
        for j in range(5, i, -1):
            if ddd[j] < 0:
                m = min(-ddd[j], ddd[i])
                ddd[j] += m
                ddd[j - i] -= m
                ddd[i] -= m
        if ddd[i] > 0:
            print('No')
            break
    else:
        print('Yes')

H - Random Kth Max

• $N$ 個の確率変数がある
• $i$ 番目は $[L_i,R_i]$ の範囲を取る連続一様分布
• 大きい方から $K$ 番目の値の期待値を求めよ
• $1 \le K \le N \le 50$
• $0 \le L_i \lt R_i \le 100$
• 入力は全て整数

$[0,2]$ と $[1,3]$ など、区間がズレながら入り交じるとややこしい。

幸い、入力は全て整数で、またその範囲も上限 $100$ と小さいので、 「$K$ 番目に大きい数が $[x,x+1)$ に入る確率を計算して、そこから期待値を $x$ ごとに計算」という手法を採れる。

これなら仮に $[x,x+1)$ の中に複数の変数が入るとしても、全ての変数はその区間をフルに動くので、区別せず扱える。

さて、ここで重要な事実として、以下がある。

• $[0,1]$ をとる連続一様分布が $n$ 個あるとき、その $i$ 番目に小さい値の期待値は $\dfrac{i}{n+1}$

たとえば $n=6$ なら、小さい方から期待値は順番に $\frac{1}{7}～\frac{6}{7}$ となる。
証明は、なんか累積確率分布を積分してごにょごにょやるとそうなるらしい。

従って、$[x,x+1)$ に何個の変数があるかで期待値が変わってくる。
さらに $[x,x+1)$ に $K$ 番目があるかどうかを判断するには $x+1$ 以上の変数の個数も必要なため、以下の情報を持ったDPが考えられる。

• $DP[x][i][j][k]=i$ 番目の区間までを考慮したとき、$[x,x+1)$ の中に入る変数が $j$ 個、$x+1$ 以上となる変数が $k$ 個となる確率

怒濤の4重ループだが、$O(R_{max} N^3)$ で $10^7$ 程度なので、 まぁPythonは厳しいかもだがPyPyやNumbaを使えば大丈夫。

実際は、$x$ については他の $x-1$ や $x+1$ に干渉することはないので $x=0,1,2,...$ と最初に決め打つことを $R_{max}$ 回やればよいし、 $i$ も直前 $i-1$ の情報があればよくinplaceに更新できるので、 実装上は $j,k$ の二次元配列でよい。

$i=N$ まで埋め終わった後、$[x,x+1)$ に $K$ 番目が入る（$k \lt K \le j+k$ となる）$j,k$ の組について、以下を求めて足し合わせれば答えとなる。

• ($[x,x+1)$ 内に $j$ 個ある時に大きい方から $K-k$ 番目の値の期待値) $\times DP[j][k]$
• $=(x+1-\dfrac{K-k}{j+1}) DP[j][k]$

DPの遷移は、$[L_i,R_i]$ の確率変数を加えるとき、 （便宜的に $i$ の次元を補って説明する）

• $0 \le x \lt L_i$ の範囲
  • 必ず $x+1$ 以上となる変数が1個増える
  • $k$ の次元をシフトすればよい
    • $DP[i+1][j][k+1]=DP[i][j][k]$ （全ての $j,k$ について）
    • $DP[i+1][j][0]=0$

• $R_i \le x$ の範囲
  • 何もしなくてよい

• $L_i \le x \lt R_i$ の範囲
• $[L_i,x)$ の値を取る
  • $DP$ の添字は変化しないが、他に遷移する分だけ、自己に遷移する分が減る
  • $DP[i+1][j][k]+=\dfrac{x-L_i}{R_i-L_i}DP[i][j][k]$
• $[x,x+1)$ の値を取る
  • $j+1$ へ遷移する
  • $DP[i+1][j+1][k]+=\dfrac{1}{R_i-L_i}DP[i][j][k]$
• $[x+1,R_i)$ の値を取る
  • $k+1$ へ遷移する
  • $DP[i+1][j][k+1]+=\dfrac{R_i-x-1}{R_i-L_i}DP[i][j][k]$

このようになる。

Python3

import os
import sys

import numpy as np


def solve(inp):
    def mod_pow(x, a, MOD):
        ret = 1
        cur = x
        while a:
            if a & 1:
                ret = ret * cur % MOD
            cur = cur * cur % MOD
            a >>= 1
        return ret

    n = inp[0]
    m = inp[1]
    lll = inp[2::2]
    rrr = inp[3::2]
    MOD = 998244353

    min_l = lll.min()
    max_r = rrr.max()

    dp = np.zeros((max_r, n + 1, n + 1), np.int64)
    # dp[x][x～x+1の変数の個数][x+1以上となる変数の個数] = 確率
    dp[:, 0, 0] = 1

    invs = np.zeros(101, np.int64)
    for i in range(1, 101):
        invs[i] = mod_pow(i, MOD - 2, MOD)

    for i in range(n):
        l = lll[i]
        r = rrr[i]
        for x in range(min_l, l):
            for j in range(i + 1):
                dp[x, j] = np.roll(dp[x, j], 1)

        for x in range(l, r):
            for j in range(i, -1, -1):
                for k in range(i - j + 1, -1, -1):
                    if dp[x, j, k] == 0:
                        continue
                    dp[x, j + 1, k] += dp[x, j, k] * invs[r - l] % MOD
                    dp[x, j + 1, k] %= MOD
                    dp[x, j, k + 1] += dp[x, j, k] * invs[r - l] % MOD * (r - x - 1) % MOD
                    dp[x, j, k + 1] %= MOD
                    dp[x, j, k] *= invs[r - l] * (x - l) % MOD
                    dp[x, j, k] %= MOD

    ans = 0
    for x in range(min_l, max_r):
        for j in range(1, n + 1):
            for k in range(m - j, m):
                if dp[x, j, k] == 0:
                    continue
                bwr = m - k
                ans += (x + 1 - bwr * invs[j + 1]) % MOD * dp[x, j, k]
                ans %= MOD

    return ans


SIGNATURE = '(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)

inp = np.fromstring(sys.stdin.read(), dtype=np.int64, sep=' ')
ans = solve(inp)
print(ans)