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

AtCoder Beginner Contest 386

E - Maximize XOR

• 長さ $N$ の非負整数列 $A$ および整数 $K$ が与えられます。ここで二項係数 $\dbinom{N}{K}$ は $10^6$ 以下であることが保証されます。
• $A$ から異なる $K$ 項を選ぶとき、選んだ $K$ 個の数の総 XOR としてあり得る最大値を求めてください。
• すなわち $\underset{1\leq i_1\lt i_2\lt \ldots\lt i_K\leq N}{\max} A_{i_1}\oplus A_{i_2}\oplus \ldots \oplus A_{i_K}$ を求めてください。

• $ 1\leq K\leq N\leq 2\times 10^5$
• $ 0\leq A_i\lt 2^{60}$
• $\dbinom{N}{K}\leq 10^6$
• 入力は全て整数

$N,K$ は大きいが、二項係数の値に制約があるという珍しい問題。
$K$ の（直接的な）制約に騙されず、実はごく小さい値しか取れないことを見抜ければ、全探索が間に合うと分かる。

二項係数の意味するところの1つである「$N$ 個の中から $K$ 個選ぶ場合の数」の、 この「選んだ具体的な $K$ 個」の組が $10^6$ 通り以下であると言っている。

よって、「$K$ 個組を全列挙する」「1組当たり $O(K)$ で実際にXORを計算する」が 現実的な時間でできるなら、実際に全列挙して最大値を求めればよい。

前者の計算量は難しいので保留して、 後者は $O(K \dbinom{N}{K})$ が実行制限時間に収まるためにはせいぜい $K \le 20$ くらいまでが限度である。

ここで、$\dbinom{N}{K} \le 10^6$ の範囲で $K$ が大きくなるケースを考えてみる。

$\dbinom{200000}{199999}$ みたいに $N,K$ がともに大きいケースは、 選ばない方を列挙することで $\dbinom{200000}{1}$ として扱える（$K$ を $N-K$ にできる）。
$K \le N-K$ を前提とすると $\dbinom{N}{K}$ が小さくなるのは $\dbinom{2K}{K}$ のような形であるが、 それでも小さい方から確認していくと $K=12$ で $\dbinom{2K}{K} \gt 10^6$ となり超えてしまう。
よって、実は $K \le 11$ が保証されていることが分かる。

前者の計算量はちゃんとはわからんけど、まぁ $K \le 11$ ならそんなひどいことにはならんやろ！と信じて投げると通る。

前者、$\dbinom{N}{K}$ の全列挙は、再帰関数によるシンプルな実装で $O(K \dbinom{N}{K})$ でできるらしい。

• 参考: QCFium氏のユーザー解説

ただ、終了条件をサボる実装だとTLEするらしく、なかなか難しい。

Pythonなら、itertools.combinations を使うのが楽。

Python3

from itertools import combinations


def solve(n, k, aaa):
    if k + k > n:
        k = n - k
        offset = 0
        for a in aaa:
            offset ^= a
    else:
        offset = 0

    ans = 0
    for chosen in combinations(aaa, k):
        tmp = offset
        for a in chosen:
            tmp ^= a
        ans = max(ans, tmp)

    return ans


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

F - Operate K

• 文字列 $S$ に対して以下の操作を $0$ 回以上 $K$ 回以下行って、文字列 $T$ と一致させられるか判定してください。
  • 次の $3$ 種類の操作のうちひとつを選択し、実行する。
    • $S$ 中の (先頭や末尾を含む) 任意の位置に、任意の文字を $1$ つ挿入する。
    • $S$ 中の文字を $1$ つ選び、削除する。
    • $S$ 中の文字を $1$ つ選び、別の $1$ つの文字に変更する。

• $S,T$ は英小文字からなる長さ $1$ 以上 $500000$ 以下の文字列
• $K$ は $\color{red}{1 \le K \le 20}$ を満たす整数

操作はレーベンシュタイン距離の操作と同じである。
2つの文字列 $S,T$ のレーベンシュタイン距離は $O(|S||T|)$ で求められる。
今回は当然それでは間に合わないが、代わりに「$K$ 以下か、より大きいか」だけ判定できればよく、$K$ は小さい。

• 解説 - AtCoder Beginner Contest 386

ここで、以下で定義されるDPのうち、

• $\mathrm{DP}[i,j]:=S$ の先頭 $i$ 文字と $T$ の先頭 $j$ 文字の距離

$|i-j| \gt K$ になるような箇所は明らかに距離が $K$ 以上となるため、埋める必要は無い。

$|i-j| \le K$ となる箇所だけ埋めていけば、$O(K(|S|+|T|))$ で求められる。

G - Many MST

• 正整数 $N,M$ が与えられます。
• 頂点に $1$ から $N$ までの番号がつけられている $N$ 頂点の重み付き完全グラフであって各辺の重みが $1$ 以上 $M$ 以下の整数であるようなものは $M^{N(N-1)/2}$ 通りあります。
• それぞれについて、最小全域木に含まれる辺の重みの和を求め、それらの総和を $998244353$ で割った余りを出力してください。

• $ 2\leq N\leq 500$
• $ 1\leq M\leq 500$
• 入力は全て整数

問題の言い換えを重ねてDPに持っていく。ある程度は典型ではあるが、言い換えの連続が難しい。

最小全域木のコストは、プリム法やクラスカル法のように、小さい辺から採用していくアルゴリズムが簡単だが、 今回は複数のグラフの最小全域木を同時に求めないといけないため、使いづらい。

他の言い換え方法として、以下がある。

• 辺重みが $0～m-1$ のグラフ $G$ を考える。
  • $k=1,...,m$ に対して「重みが $k$ 未満の辺のみを残したグラフ」を $G_k$ とする
  • その連結成分の個数を $c(G_k)$ とする
  • 最小全域木のコストは、$\displaystyle \sum_{k=1}^{m}(c(G_k)-1)$

辺の重みを1階層ずつに分解するイメージ。
例えば重み $5$ の辺が最小全域木に使われるとして、 それによって結ばれる2つの連結成分を考えると、 $k=1～5$ に対しては結ばれないので連結成分数は各「1」だけ多くなり、 それ以上に対しては結ばれる。合計「5」だけ全体に寄与することになる。

つまり、$M^{N(N-1)/2}$ 個ある全てのグラフの集合を $S$ として、 $\displaystyle \sum_{G \in S}\sum_{k=1}^{m}(c(G_k)-1)$ が答えとなる。

$\displaystyle \sum_{G \in S}\sum_{k=1}^{m}(c(G_k)-1)$ の定数部分を外に出すと、 $\displaystyle \sum_{G \in S}\sum_{k=1}^{m}c(G_k) - M \times M^{N(N-1)/2}$ となる。

最終的に、$\displaystyle \sum_{G \in S}\sum_{k=1}^{m}c(G_k) + (N-M-1) \times M^{N(N-1)/2}$ を求めればよい。

後半は簡単に求まるので、$\displaystyle \sum_{G \in S}\sum_{k=1}^{m}c(G_k)$ を求められればよい。

Σの順番を入れ替え、$k$ 毎に $\sum_{G \in S}c(G_k)$ を考えることにする。

主客転倒し、「ある特定の頂点集合 $T$ が、$G_k$ において1つの連結成分を構成する」 （そして答えに1寄与する）ような $G \in S$ がいくつあるか？を考える。
$2^N-1$ 個ある全ての頂点集合について、それを $T$ としたときの結果を合計すれば、$\sum_{G \in S}c(G_k)$ が求まる。

$|T|$ が同じなら求めたい値も一緒なので、集合のサイズだけが重要となる。
$s=1,2,...,N$ 毎に「$f_k(s):=$ 頂点集合 $\{1,2,...,s\}$ が $G_k$ において1つの連結成分となる $G$ の個数」を求めれば、その結果をそれぞれ $\dbinom{N}{s}$ 倍すればよい。

で、その求め方だが、

(a)     (b)    (c)
○--○        ○  ○    (a): Tに属する頂点同士を結ぶ辺を表す
｜×｜ <---->   ○      (b): TとT以外を結ぶ辺を表す
○--○        ○  ○    (c): T以外同士を結ぶ辺を表す

• (a) $s$ 個の頂点をコスト $k$ 未満の辺で連結にしつつ、残りの辺のコストも決める方法の個数
  • ちょっと難しいので、とりあえず $g_k(s)$ と置く
• (b) 特定の $s$ 個と残りの $N-s$ 個を結ぶ $s(N-s)$ 本の辺のコストの決め方の個数
  • $G_k$ において非連結なので、全てコスト $k$ 以上でないといけない
• (c) 残りの $N-s$ 個同士を結ぶ辺のコストの決め方の個数
  • これは自由

$g_k(s)$ が求まれば、$f_k(s) = g_k(s) \times (M-k)^{s(N-s)} \times M^{(N-s)(N-s-1)/2}$ で求められる。

$g_k(s)$ を $s$ の小さい方から求める。
まず、$g_k(1)=1$ である。

$s \ge 2$ に対しては、

• (a) 内の全ての辺を自由に決める方法の個数は $M^{s(s-1)/2}$
• そのうち、連結成分が2つ以上に分かれるような方法の個数を引く。
  • 頂点 $1$ を含む連結成分のサイズが $t$（$t=1,...,s-1$）となるような決め方の個数は、上記の(a)(b)(c)と同様の考え方により、$\dbinom{s-1}{t-1} \times g_k(t) \times (M-k)^{t(s-t)} \times M^{(s-t)(s-t-1)/2}$
• $\displaystyle g_k(s) = M^{s(s-1)/2} - \sum_{t=1}^{s-1} \dbinom{s-1}{t-1} \times g_k(t) \times (M-k)^{t(s-t)} \times M^{(s-t)(s-t-1)/2}$

これで、$O(N^2)$ で $g_k$ を求められる。

各 $k$ に対して繰り返し、$O(N^2M)$ で全て求められる。

Python3

import numpy as np
from numba import njit


@njit('(i8,i8)', cache=True)
def solve(n, m):
    MOD = 998244353

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

    pow_lim = n * (n - 1) // 2 + 1
    pows = np.ones((m + 1, pow_lim), dtype=np.int64)
    pows[0, 1:] = 0
    for k in range(2, m + 1):
        for i in range(1, pow_lim):
            pows[k, i] = pows[k, i - 1] * k % MOD

    def get_connected_with_under_k_edges(k):
        """
        i=1,2,...,N 頂点の、各辺に 0～M-1 の値が付いた完全グラフで、
        k 未満の辺によって全体が連結となっているものの、値の割り振り方を含めた個数
        """
        dp = np.zeros(n + 1, np.int64)  # i 頂点 j 辺の連結グラフの個数
        dp[1] = 1
        for i in range(2, n + 1):
            res = pows[m, i * (i - 1) // 2]
            # 連結でないものを除く。頂点1が含まれる連結成分のサイズを j に固定する。
            for j in range(1, i):
                tmp = dp[j]
                tmp *= facts[i - 1] * finvs[j - 1] % MOD * finvs[i - j] % MOD
                tmp %= MOD
                tmp *= pows[m - k, j * (i - j)]
                tmp %= MOD
                tmp *= pows[m, (i - j) * (i - j - 1) // 2]
                tmp %= MOD
                res -= tmp
                res %= MOD
            dp[i] = res
        return dp

    ans = pows[m, n * (n - 1) // 2] * (n - m - 1) % MOD

    for k in range(1, m + 1):
        res = get_connected_with_under_k_edges(k)
        for i in range(1, n + 1):
            tmp = res[i]
            tmp *= facts[n] * finvs[i] % MOD * finvs[n - i] % MOD
            tmp %= MOD
            tmp *= pows[m - k, i * (n - i)]
            tmp %= MOD
            tmp *= pows[m, (n - i) * (n - i - 1) // 2]
            tmp %= MOD
            ans += tmp
            ans %= MOD

    return ans


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