AtCoder Beginner Contest 243 E,F,G問題メモ

AtCoder Beginner Contest 243

4完だった。ぐぬぬ。

E - Edge Deletion

• $N$ 頂点 $M$ 辺の重み付き単純無向連結グラフが与えられる
  • $A_i,B_i,C_i$: $A_i$ と $B_i$ を相互に結ぶ辺の重みは $C_i$
• 「削除してもどの2頂点の最短距離も変化しない」辺を最大何本削除できるか、求めよ
• $2 \le N \le 300$

各辺につき、迂回路があるかどうか、という問題。

    5       ①-②を結ぶ重み5の辺があるが、
①----②    その他の辺を使って①→②に距離5でたどり着けるので、
2＼  ／3    この辺は不要
   ③

なので、「必ず2本以上の辺を使う場合の全頂点間最短距離 $Dist(u,v)$」を計算してやって、
各辺につき、$Dist(A_i,B_i) \le C_i$ なら、その辺は不要と判断できる。

$Dist(u,v)$ は、普通にFloyd-Warshallで（1本の辺だけでもよい）全頂点最短距離 $d(u,v)$ を求めた後、
必要な $(u,v)$ について、$\displaystyle \min_{w \neq u,v}(d(u,w)+d(w,v))$ を求めればよい。

①-②の辺は、①-③、③-②が残るなら確かに不要そうだけど、これらの辺も削除されちゃうかもしれないし、 全体的にはやっぱり①-②は必要だった、ということにはならないの？

ならない。
この方針に従う限り、たとえば①-③の辺が削除できるなら、 ①-③間をそれ以下のコストで代替できる迂回路が存在するので、①-②の迂回路の一部としてもそっちを使えばよい。

きちんとした証明としては、公式解説の通り、

• その1本だけの削除であっても、削除したら絶対に最短距離が変わる頂点がある辺、という条件を整理する
• その辺以外を全て削除すると考える（答えの上界）
• 削除後のグラフで、最短距離が変化していない（上界が達成できる）ことを、削除条件などから証明する

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))

INF = 1 << 60
dists1 = [[INF] * n for _ in range(n)]
dists2 = [[INF] * n for _ in range(n)]
for u in range(n):
    for v, c1 in links[u]:
        dists1[u][v] = c1
        for w, c2 in links[v]:
            if u != w:
                dists2[u][w] = min(dists2[u][w], c1 + c2)

for k in range(n):
    for i in range(n):
        for j in range(n):
            dists2[i][j] = min(dists2[i][j],
                               dists1[i][k] + dists2[k][j],
                               dists2[i][k] + dists1[k][j],
                               dists2[i][k] + dists2[k][j])

ans = 0
for i in range(n):
    for j, c in links[i]:
        if i < j and dists2[i][j] <= c:
            ans += 1

print(ans)

F - Lottery

• $N$ 種類の景品が当たるくじがある
  • 当選の重み $W_1,W_2,...,W_N$ が決まっており、$k$ 番目の景品の抽選率は $\dfrac{W_k}{W_iの総和}$
• くじを $K$ 回引いたとき（抽選率は毎回独立）、ちょうど $M$ 種類の景品を得る確率を $998244353$ の有理数modで求めよ
• $1 \le M \le N \le 50$
• $1 \le K \le 50$
• $1 \le W_i \lt 998244353$、$W_i$ の総和もmodの値未満

ちょうど $M$ 種類、というのはややこしくても、$M$ 種類以下、というのはまだ求めやすい。

これをまず求めて、そこから包除原理かなんかでちょうどでないものを除く、 というのがよくある解法なので、この問題もそれで考える。

全 $W_i$ の総和を $\displaystyle S = \sum_{i=1}^N W_i$ とする。

ある $m$ 個の選び方を固定（$m_j$）し、その $W_i$ の和を $P_{mj}$ と表すとする。

m=3 で m1={1,3,4} という選び方を固定すると、Pm1 = W1 + W3 + W4

同じ m 個でも、違う選び方 m2,m3,... でそれぞれ Pm2,Pm3,... は異なる

すると、「くじ $K$ 回を通して $m_1$ の中のどれかしか当たらない確率」は、$\left ( \dfrac{P_{m1}}{S} \right )^K$ である。

ここで、全ての項に $\dfrac{1}{S^K}$ は共通なので最後にまとめてかければよく、分子のみを考えることにする。

まずは、$m=M$ の全選び方 $M_1,M_2,...$ に対して、 $\displaystyle f(m) = \sum_{j} P_{mj}^K$（つまり、全 $P_{mj}$ の $K$ 乗の総和）を求めたい。

DPで、$W_i$ を順番に加えていくことを考える。

• $DP[i][m][k]=W_i$ までを考慮し、そのうち $m$ 個の要素からなる全選び方に対し、$P_{mj}^k$ の総和

$f(M)=DP[N][M][K]$ である。

$DP[i-1][m][k]=P_{m1}^k+P_{m2}^k+...$ が求まっているとして、ここに新たに $W_i$ を加えた選び方を考慮する。

• $DP[i][m+1][k] += (P_{m1} + W_i)^k + (P_{m2} + W_i)^k + ...$

これを展開整理すると、（長くなるので、$DP[i-1][m][k]=d_k$ とおくと）

• $DP[i][m+1][k] += d_k + {}_kC_1 d_{k-1} W_i + {}_kC_2 d_{k-2} W_i^2 + ... + {}_kC_{k-1} d_1 W_i^{k-1} + W_i^k$

なので、これまでに求めた成果から、1回の遷移あたり $O(k)$ で求められる。

全部で $O(NMK^2)$ で、このDPを埋められる。

$f(M)=DP[N][M][K]$ には、ちょうど $M$ 種類でないものも含まれる。

ここから重複を除いた、「ちょうど」$M$ 種類である求めるべき答えを $g(M)$ とする。

たとえば $N=5,M=4$ での選び方に対し、

    1 2 3 4 5
M1  o o o o
M2  o o o   o
M3  o o   o o
M4  o   o o o
M5    o o o o

• $\{1,2,3\}$ などの $M-1$ 個の選び方は、$M_1$ と $M_2$ で2回分
• $\{1,2\}$ などの $M-2$ 個の選び方は3回分
• $\{1\}$ のような $M-3$ 個の選び方は4回分

このように「余分な状態が、しかも重複して」数えられている。

ただし、それぞれいくつ重複しているかは、集合のサイズごとに固定 （$\{1,2,3\}$ も $\{2,3,4\}$ も重複している個数は一緒）なので、まとめて除くことが可能である。

$M-i$ 種類しか選ばれない状態が $f(M)$ において重複している回数は、 残りの $i$ 種類を、選ばれていない $N-(M-i)$ 個から選ぶ ${}_{N-(M-i)}C_{i}$ である。

よって、$i=1～M-1$ に対し、$f(M) -= {}_{N-(M-i)}C_{i} g(M-i)$ をすればよい。最終的に残るのが $g(M)$ である。

途中に出てくる $g(m)$ については、$m$ が小さい方から順に求めていけばよい。$O(M^2)$ で計算可能。

Python3

import os
import sys

import numpy as np


def solve(inp):
    n, m, l = inp[:3]
    www = inp[3:]
    MOD = 998244353

    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

    def prepare_factorials(n, MOD):
        factorials = np.ones(n + 1, np.int64)
        for m in range(1, n + 1):
            factorials[m] = factorials[m - 1] * m % MOD
        inversions = np.ones(n + 1, np.int64)
        inversions[n] = mod_pow(factorials[n], MOD - 2, MOD)
        for m in range(n, 1, -1):
            inversions[m - 1] = inversions[m] * m % MOD
        return factorials, inversions

    facts, finvs = prepare_factorials(n + l, MOD)

    def ncr(n, r):
        return facts[n] * finvs[r] % MOD * finvs[n - r] % MOD

    dp = np.zeros((m + 1, l + 1), np.int64)
    for i in range(n):
        w = www[i]
        for j in range(min(i, m - 1), 0, -1):
            for k in range(l + 1):
                tmp = 0
                for e in range(k + 1):
                    tmp += dp[j, k - e] * ncr(k, e) % MOD * mod_pow(w, e, MOD) % MOD
                dp[j + 1, k] += tmp % MOD
        for k in range(l + 1):
            dp[1, k] += mod_pow(w, k, MOD)
        dp %= MOD

    for i in range(2, m + 1):
        for j in range(1, i):
            dp[i, l] -= dp[j, l] * ncr(n - j, i - j) % MOD
            dp[i, l] %= MOD

    ans = dp[m, l] * mod_pow(mod_pow(www.sum(), l, MOD), MOD - 2, MOD) % 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)

G - Sqrt

• はじめ、1項だけの整数列 $A=[X]$ がある
• $A$ の末尾の項を $x$ として、$1～floor(\sqrt{x})$ の整数を好きに選び追加することを $10^{100}$ 回繰り返す
• 最終的にできうる整数列は何通りあるか求めよ
• $T$ ケース与えられるので、それぞれについて求めよ
• $1 \le T \le 20$
• $1 \le X \le 9 \times 10^{18}$

繰返し回数がとんでもない数だが、末尾が $1$ になったら以降、$1$ を繰り返すしかないし、 $10^{100}$ 回までには絶対 $1$ になるので、「$1$ になったら終了」と考えてよい。

以下、簡単のため $\sqrt{x}$ だけで床関数適用済みの $floor(\sqrt{x})$ を表すとする。

$X=x$ の答えを $f(x)$ とする。
たとえば $x=30$ だったら、その直後には 5～1 を置けるので、$f(5)+f(4)+f(3)+f(2)+f(1)$ が答えとなる。

一般化すると、$f(x)=f(1)+f(2)+...+f(\sqrt{x})$ である。

$f(1)=1$ からはじめて $f(2),f(3),...$ と1つずつ求めていけば求まるのは求まるが、 $X$ が馬鹿でかいので、平方根を取ったとしても最大 $f(3 \times 10^9)$ を求めなければならない。時間的に無理。

$x$ を順番に書き出すと、かなり同じ値が連続することがわかる。

  x   1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 ..
f(x)  1  1  1  2  2  2  2  2  3  3  3  3  3  3  3  5  5 ..

変曲点は、$\sqrt{x}$ が変化する点、つまり $x$ が平方数になったときにあるので、そこを中心にまとめると、

  x   1 ....  4 .....  9 ..... 16 ..... 25 ......  36 .... 49 .. 64 .. 81 .. 100 ..
f(x)  1 .. 1  2 ... 2  3 ..  3  5 ... 5  7 .... 7   9 .... 11 .. 13 .. 16 ..  19 ..
acc   1 .. 3  5 .. 13 16 .. 34 39 .. 79 86 .. 156 165 ..

accは、$f(x)$ の累積和であり、$f(x)=acc(\sqrt{x})$ となる。つまりは $acc(\sqrt{X})$ を求めたい。

• A097046 - OEIS
  • accはこの数列になるが、一般項を求める式はさすがに無さそう

とりあえず、$f(x)$ が同じ区分ごとにまとめなおそう。
添字 $i$ に対応する $f(x)$ を $g(i)$ とする。
cntはその区分の長さで、隣接する平方数の差なので要は奇数列となる。

  x       1 ....  4 .....  9 ..... 16 ..... 25 ......  36 .... 49 .. 64 .. 81 .. 100 ..
f(x)      1 .. 1  2 ... 2  3 ..  3  5 ... 5  7 .... 7   9 .... 11 .. 13 .. 16 ..  19 ..
acc       1 .. 3  5 .. 13 16 .. 34 39 .. 79 86 .. 156 165 ..

  i    (0)     1        2        3        4         5        6     7     8     9
g(i)           1        2        3        5         7        9    11    13    16
cnt            3        5        7        9        11       13    15    17    19
acc末尾 0      3       13       34       79       156      273   438   658   963

こうすると、添字 $i=\sqrt{x}$ を表すことになり、 $i$ に対して求めるべき範囲は最初から比べると $\sqrt[4]{X}$ まで減少する。
これなら1つずつ求めても間に合う。

$g(i)$ を求めるとき、$j=\sqrt{i}$、端数を $k$ とする（$i=j^2+k$ となる）と、

• $g(i)=acc末尾[j-1]+(k+1)g(j)$
• $acc末尾[i] = acc末尾[i-1] + g(i)cnt[i]$

acc末尾で直前 $j-1$ までの累積和を求めて、そこから $k+1$ 進んだ数、という感じで計算できる。

$g(i),acc末尾[i]$ の2つを $i=10^5$ くらいまで求めておけば、クエリに対して $O(1)$ で計算できる。

4乗根 $j=\sqrt{\sqrt{X}}$、端数 $k=\sqrt{X}-j^2$ とすると、以下が答えとなる。

• $acc末尾[j-1]+(k+1)g(j)$

Pythonだと、制約上限から1小さい数 x=8999999999999999999 の平方根を int(x ** 0.5) で求めると、3000000000 になってしまった。 整数から浮動小数にする段階で53bitの精度になるので、64bit型の上限付近だと狂いが生じるので注意。

Python3

from decimal import Decimal, ROUND_DOWN

lim = 10 ** 5
ans = [0, 1]
acc = [0, 3]
for i in range(2, lim):
    odd = 2 * i + 1
    sq = int(i ** 0.5)
    md = i - sq ** 2
    ans.append(acc[sq - 1] + ans[sq] * (md + 1))
    acc.append(acc[-1] + odd * ans[i])

t = int(input())
qua = Decimal('1')

for _ in range(t):
    x4 = Decimal(input())
    x2 = x4.sqrt().quantize(qua, rounding=ROUND_DOWN)
    x = x2.sqrt().quantize(qua, rounding=ROUND_DOWN)
    md = int(x2 - x ** 2)
    print(acc[int(x) - 1] + ans[int(x)] * (md + 1))