AtCoder Beginner Contest 352 F,G問題メモ

AtCoder Beginner Contest 352

F - Estimate Order

問題

$N$ 人の人が競争して、 $1～N$ の順位が付いた（同率無し）
以下のような $M$ 個の条件が与えられる
- （ $A_i$ の人の順位）-（ $B_i$ の人の順位）＝ $C_i$ である
- 矛盾するような条件は与えられない
人 $1,...,N$ について、「与えられた条件だけで順位が一意に定まる」場合はその順位を、定まらない場合は '-1' を答えよ
$2 \le N \le 16$

解法

やること自体は手計算でも十分にパズルとして解けるような内容だが、実装にちょっと迷走してしまった。

与えられた条件をグラフのように結ぶと、連結成分ごとに1つの「パズルのピース」ができる。

Ai  Bi  Ci
 1   2   3        ①・・②・④・③  条件により結びついた人達の間で、
 2   3   4    →  |--3--||---4---|  この位置関係は変わらない
 1   4   5        |----5----|

              →  1 0 0 1 0 1 0 1   bitフラグで表しておく
              
他の連結成分からも、  1  や  101  などのピースができる

で、各人の順位を求めるのは、長さ $N$ の 0000…00 という“枠”にこれらのピースを、1 が重ならないように全て当てはめるという操作になる。
制約上、必ず1通りは正しい当てはめ方がある。

一意に決まるかどうかは、以下のように確かめる。

0000000000  長さNの枠

10010101    ピースAの1つを試しに置く

 1    1001  残りのピースの置き方を全探索し、
  101   1   残りを全て埋める方法があるか？をチェック
            →埋められた（1つ見つかれば、即終了してよい）

 10010101   ピースAの位置をずらして、
  10010101  他の場合も、残りを全て埋められるかチェック

で、「残りを全て埋められるピースAの初期位置」が1箇所だけだったら、ピースAに属する人の順位は全て決まる。

2箇所以上「残りを全て埋められる、ピースAの初期位置」があったら定まらない。

$N$ が小さいので、全てのピースについて、この全探索を実行しても間に合う。

Python3

        
          
          
              
              n, m = map(int, input().split())
links = [[] for _ in range(n)]
for _ in range(m):
    a, b, c = map(int, input().split())
    a -= 1
    b -= 1
    links[a].append((b, -c))
    links[b].append((a, c))
 
groups = []
checked = [False] * n
for s in range(n):
    if checked[s]:
        continue
    q = [s]
    relations = {s: 0}
    while q:
        a = q.pop()
        for b, c in links[a]:
            if b in relations:
                continue
            relations[b] = relations[a] + c
            checked[b] = True
            q.append(b)
    offset = min(relations.values())
    use_rank = 0
    for a in relations:
        relations[a] -= offset
        use_rank |= 1 << relations[a]
    groups.append([relations, use_rank])
 
 
def check(groups, checking, i, remaining_j, memo={}):
    if (i, remaining_j) in memo:
        return memo[i, remaining_j]
    if i == len(groups):
        return True
    if i == checking or fix_group[i]:
        return check(groups, checking, i + 1, remaining_j)
    use_rank = groups[i][1]
    while use_rank <= remaining_j:
        if use_rank & remaining_j == use_rank:
            res = check(groups, checking, i + 1, remaining_j ^ use_rank)
            if res:
                memo[i, remaining_j] = True
                return True
        use_rank <<= 1
    memo[i, remaining_j] = False
    return False
 
 
ans = [-1] * n
remaining_j = (1 << n) - 1
len_grp = len(groups)
fix_group = [False] * len_grp
while True:
 
    updated = False
    for gi in range(len_grp):
        if fix_group[gi]:
            continue
        relations, use_rank = groups[gi]
        ok_rank = -1
        while use_rank <= remaining_j:
            if use_rank & remaining_j == use_rank:
                res = check(groups, gi, 0, remaining_j ^ use_rank)
                if res:
                    if ok_rank == -1:
                        ok_rank = use_rank
                    else:
                        ok_rank = -1
                        break
            use_rank <<= 1
 
        # 他と矛盾しないずらし方が1通りしかない
        if ok_rank != -1:
            lsb = (ok_rank & -ok_rank).bit_length() - 1
            for i, j in relations.items():
                ans[i] = j + lsb + 1
            fix_group[gi] = True
            remaining_j ^= ok_rank
            updated = True
            break
 
    if not updated:
        break
 
print(*ans)

            

        
      

G - Socks 3

問題

タンスに靴下がいっぱい入っている。種類 $i$ の靴下は $A_i$ 枚入っている
今日はく靴下を取り出すために、以下の操作を行う
- ランダムに等確率で靴下を1枚取り出すことを繰り返す（一度取り出した靴下は戻さない）
- 同じ種類の靴下が2枚取り出された時点で、取り出すのをやめる
靴下を取り出す回数の期待値を求めよ
$1 \le N \le 3 \times 10^5$
$2 \le A_i \le 3000$

解法

このように操作回数の期待値を求める問題において、「ちょうど」の期待値は、「○○回以上」の確率の和に置き換えられる。

主客転倒の一種。

ゲームの確率遷移の例（この問題とは関係ない）

        1回      2回     3回 ...
1  →  1/4 --→ 1/8[終]
 |        `--→ 1/8[終]
 |-→  1/2 --→ 1/3 --→ 1/6[終]
 |        |       `---→ 1/6[終]
 |        `--→ 1/6[終]
 `-→  1/4[終]
 
 操作回数の期待値 ＝ [終]のノードが表す回数 × そのノードに至る確率 の総和
   1*1/4 + 2*(1/8+1/8+1/6) + 3*(1/6+1/6)
 
 これは、[終]に限らない、全てのノードの確率の合計と一致する
   1/4 + 1/8 + 1/8 + 1/2 + 1/3 + 1/6 + 1/6 + 1/6 + 1/4

元の問題に戻る。

$A_i$ の総和を $S$ とする。

たとえば「3枚」取り出してもまだ揃っていない場合というのは、全て別々の種類から取り出すということなので、

取り出す靴下の種類と順序を固定（例） $A_1→A_4→A_3$
1回目、 $\dfrac{A_1}{S}$
2回目、 $\dfrac{A_4}{S-1}$
3回目、 $\dfrac{A_3}{S-2}$
→ $\dfrac{A_1A_4A_3}{S(S-1)(S-2)}$

というのを、全ての靴下の種類・順序について足し合わせたものとなる。

これは形式的冪級数を用いて

$F(x)=(1+A_1x)(1+A_2x)(1+A_3x)...(1+A_Nx)$

とすると、 $x^3$ の係数が、相異なる全ての組 $(i,j,k)$ についての $A_iA_jA_k$ を足し合わせたものとなる。
（同様に、 $x^p$ の係数が、全ての $p$ 個組についての $A_iA_jA_k...$ を足し合わせたものとなる）

順序を考慮するための階乗をかけて、以下を足し合わせたものが、操作が終わっていない状態全ての確率の和となる。

$\dfrac{[x]F(x) \times 1!}{S}$
$\dfrac{[x^2]F(x) \times 2!}{S(S-1)}$
$\dfrac{[x^3]F(x) \times 3!}{S(S-1)(S-2)}$
…
$\dfrac{[x^N]F(x) \times N!}{S(S-1)(S-2)...(S-N+1)}$

最後に、操作終了のための1枚を加えて、答えとなる。

$F(x)$ を求める際はFFTを用いるが、二分木の要領で、隣り合う要素同士を順次マージしてなるべく長い配列を何度もマージしないように工夫すると、全体を $O(N \log^2{N})$ で求められる。

Python3

な～んか、自前のFFTライブラリ遅いんだよな……

        
          
          
              
              import os
import sys
 
from collections import deque
 
import numpy as np
 
input = sys.stdin.buffer.readline
 
# MOD 998244353 前提
MOD = 998244353
SUM_E = np.array(
    [911660635, 509520358, 369330050, 332049552, 983190778, 123842337, 238493703, 975955924, 603855026, 856644456,
     131300601, 842657263, 730768835, 942482514, 806263778, 151565301, 510815449, 503497456, 743006876, 741047443,
     56250497, 867605899, 0, 0, 0, 0, 0, 0, 0, 0], np.int64)
SUM_IE = np.array(
    [86583718, 372528824, 373294451, 645684063, 112220581, 692852209, 155456985, 797128860, 90816748, 860285882,
     927414960, 354738543, 109331171, 293255632, 535113200, 308540755, 121186627, 608385704, 438932459, 359477183,
     824071951, 103369235, 0, 0, 0, 0, 0, 0, 0, 0], np.int64)
 
 
def convolution(aaa, bbb):
    def bit_length(n):
        x = 0
        while 1 << x < n:
            x += 1
        return x
 
    def bit_scan_forward(n):
        x = 0
        while n & 1 == 0:
            n >>= 1
            x += 1
        return x
 
    def mod_pow(x, a, MOD):
        ret = 1
        cur = x
        while a > 0:
            if a & 1:
                ret = ret * cur % MOD
            cur = cur * cur % MOD
            a >>= 1
        return ret
 
    def butterfly(aaa, mod, sum_e):
        n = aaa.size
        h = bit_length(n)
 
        for ph in range(1, h + 1):
            w = 1 << (ph - 1)
            p = 1 << (h - ph)
            now = 1
            for s in range(w):
                offset = s << (h - ph + 1)
                for i in range(p):
                    l = aaa[i + offset]
                    r = aaa[i + offset + p] * now % mod
                    aaa[i + offset] = (l + r) % mod
                    aaa[i + offset + p] = (l - r) % mod
                now = now * sum_e[bit_scan_forward(~s)] % mod
 
    def butterfly_inv(aaa, mod, sum_ie):
        n = aaa.size
        h = bit_length(n)
 
        for ph in range(h, 0, -1):
            w = 1 << (ph - 1)
            p = 1 << (h - ph)
            inow = 1
            for s in range(w):
                offset = s << (h - ph + 1)
                for i in range(p):
                    l = aaa[i + offset]
                    r = aaa[i + offset + p]
                    aaa[i + offset] = (l + r) % mod
                    aaa[i + offset + p] = ((l - r) * inow) % mod
                inow = inow * sum_ie[bit_scan_forward(~s)] % mod
 
    n = aaa.size
    m = bbb.size
    k = n + m - 1
    z = 1 << bit_length(k)
    raaa = np.zeros(z, np.int64)
    rbbb = np.zeros(z, np.int64)
    raaa[:n] = aaa
    rbbb[:m] = bbb
    butterfly(raaa, MOD, SUM_E)
    butterfly(rbbb, MOD, SUM_E)
    for i in range(z):
        raaa[i] = raaa[i] * rbbb[i] % MOD
    butterfly_inv(raaa, MOD, SUM_IE)
    iz = mod_pow(z, MOD - 2, MOD)
    for i in range(k):
        raaa[i] = raaa[i] * iz % MOD
    return raaa[:k]
 
 
SIGNATURE = '(i8[:],i8[:])'
if sys.argv[-1] == 'ONLINE_JUDGE':
    from numba.pycc import CC
 
    cc = CC('my_module')
    cc.export('convolution', SIGNATURE)(convolution)
    cc.compile()
    exit()
 
if os.name == 'posix':
    # noinspection PyUnresolvedReferences
    from my_module import convolution
else:
    from numba import njit
 
    convolution = njit(SIGNATURE, cache=True)(convolution)
    print('compiled', file=sys.stderr)
 
n = int(input())
aaa = list(map(int, input().split()))
q = deque()
s = 0
for a in aaa:
    q.append(np.array([1, a], np.int64))
    s += a
 
while len(q) > 1:
    arr1 = q.popleft()
    arr2 = q.popleft()
    arr3 = convolution(arr1, arr2)
    q.append(arr3)
 
res = q[0]
f = 1
for i in range(2, n + 1):
    f = f * i % MOD
    res[i] = res[i] * f % MOD
 
spk = 1
for i in range(s, s - n, -1):
    spk *= i
    spk %= MOD
spk_inv = [pow(spk, -1, MOD)]
for i in range(s - n + 1, s + 1):
    spk_inv.append(spk_inv[-1] * i % MOD)
spk_inv.reverse()
 
ans = 1
for i in range(1, n + 1):
    ans += res[i] * spk_inv[i]
    ans %= MOD
 
print(ans)

            

        
      