AtCoder Beginner Contest 242 F,G問題メモ

AtCoder Beginner Contest 242

最近のABCのCとかD、普通にむずいんだよな

F - Black and White Rooks

問題

$N \times M$ マスの盤面上にブラック飛車 $B$ 個とホワイト飛車 $W$ 個を置く
配置ルール
- 1マスに2個以上の駒を置かない
- 違う色の飛車同士は互いの利きにいない（同じ色はOK）
同じ色の駒は互いに区別しないとき、置き方の個数を $998244353$ で割ったあまりで求めよ
$1 \le N,M \le 50$

解法

まず、1つの行や列に、黒と白両方が置かれることはない。

「●・・●・・・○」みたいな配置を考えたとき、混在していると必ず黒と白の境界が存在するので、その部分の駒が互いの利きに入ってしまう。

なので、行と列をそれぞれ、以下のグループに分ける。

黒のみ置かれる
白のみ置かれる
どちらも置かれない

「黒のみ置かれる行数 $R_b$」「黒のみ置かれる列数 $C_b$」「白のみ置かれる行数 $R_w$」「白のみ置かれる列数 $C_w$」としてこれらを決め打ったとき、その行・列を $N \times M$ のどこに配置するかは、以下のようになる。

$\dfrac{N!}{R_b!R_w!(N-R_b-R_w)!}\dfrac{M!}{C_b!C_w!(M-C_b-C_w)}$

また、以下がわかれば、

$f_b(R_b,C_b)=$ その行数・列数内での黒い駒 $B$ 個の置き方
$f_w(R_w,C_w)=$ その行数・列数内での白い駒 $W$ 個の置き方

全体の置き方の個数は、これらの積となる。

$\dfrac{N!}{R_b!R_w!(N-R_b-R_w)!}\dfrac{M!}{C_b!C_w!(M-C_b-C_w)} \times f_b(R_b,C_b) \times f_w(R_w,C_w)$

$f_b,f_w$ をDPで求める

$R_b,C_b,R_w,C_w$ の全探索は $O(N^2M^2)$ で可能なので、あとは $f_b,f_w$ を前計算などで $O(1)$ で得られればよい。

重複を避けるため、$f_b,f_w$ は必ず空行・空列のない置き方で無いといけない。

B=3 Rb=3 Cb=2
OK      NG      NG      1つめのNGの置き方は、
●・・  ●・・  ●●●  Rb=2 Cb=2 を探索したときに既に数えられている
・●●  ●●・  ・・・

重複を気にしなければ、$R_b \times C_b$ 個のマスの中から置く $B$ マスを選べば良いので、${}_{R_b \times C_b}C_{B}$。

そこから空になってしまう行数・列数を $R_x,C_x$ と決め打てば、除外すべき置き方はどの行・列を空にするかを含めて、以下のようになる。

${}_{R_b}C_{R_x} \times {}_{C_b}C_{C_x} \times f_b(R_b-R_x,C_b-C_x)$

これを $0 \le R_x \le R_b, 0 \le C_x \le C_b$（ただし $R_x=R_b$ かつ $C_x=C_b$ を除く）の範囲で全探索し、 ${}_{R_b \times C_b}C_{B}$ から除外すれば、$f_b(R_b,C_b)$ が求められる。
これも $O(N^2M^2)$ で計算できる。

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

    def precompute_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

    def get_dp(b, n, m, facts, finvs, MOD):
        dp = np.zeros((min(n, b) + 1, min(m, b) + 1), np.int64)
        for i in range(1, min(n, b) + 1):
            for j in range(1, min(m, b) + 1):
                if i * j < b:
                    continue
                res = facts[i * j] * finvs[b] % MOD * finvs[i * j - b] % MOD
                for rem_i in range(1, i):
                    res -= dp[i - rem_i, j] * facts[i] % MOD * finvs[rem_i] % MOD * finvs[i - rem_i] % MOD
                for rem_j in range(1, j):
                    res -= dp[i, j - rem_j] * facts[j] % MOD * finvs[rem_j] % MOD * finvs[j - rem_j] % MOD
                for rem_i in range(1, i):
                    for rem_j in range(1, j):
                        res -= dp[i - rem_i, j - rem_j] \
                               * facts[i] % MOD * finvs[rem_i] % MOD * finvs[i - rem_i] % MOD \
                               * facts[j] % MOD * finvs[rem_j] % MOD * finvs[j - rem_j] % MOD
                dp[i, j] = res % MOD
        return dp

    n, m, b, w = inp
    MOD = 998244353
    facts, finvs = precompute_factorials(n * m, MOD)

    dp_b = get_dp(b, n, m, facts, finvs, MOD)
    dp_w = get_dp(w, n, m, facts, finvs, MOD)

    ans = 0
    for rb in range(1, min(n, b + 1)):
        for cb in range(min(m - 1, b), 0, -1):
            if rb * cb < b:
                break
            for rw in range(1, min(n - rb + 1, w + 1)):
                for cw in range(min(m - cb, w), 0, -1):
                    if rw * cw < w:
                        break
                    r_pat = facts[n] * finvs[rb] % MOD * finvs[rw] % MOD * finvs[n - rb - rw] % MOD
                    c_pat = facts[m] * finvs[cb] % MOD * finvs[cw] % MOD * finvs[m - cb - cw] % MOD
                    ans += dp_b[rb, cb] * dp_w[rw, cw] % MOD * r_pat % MOD * c_pat % MOD
                    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)

包除原理で求める

何の個数を $(-1)^k$ の $k$ とすればよいか少し迷うが、「必ず空行・空列にすると定める行数+列数」とできる。

空行があってもよい条件でそれぞれの場合を求め、

$\dfrac{N!}{R_b!R_w!(N-R_b-R_w)!}\dfrac{M!}{C_b!C_w!(M-C_b-C_w)} \times {}_{R_b \times C_b}C_{B} \times {}_{R_w \times C_w}C_{W}$

その上で $(N+M-R_b-C_b-R_w-C_w)$ の偶奇によって答えに足す正負を切り替えると、最終的な答えが求まる。

これなら前計算は必要なく、$R_b,C_b,R_w,C_w$ の総当たりのみで済む。

相変わらず、包除原理はなかなか直感的には見えづらい。

G - Range Pairing Query

なんでこんな解かれてるの？

問題

長さ $N$ の数列 $A_1,...,A_N$ が与えられる
$Q$ 個のクエリを処理せよ
- $l,r$ が与えられる
- $A_{l},A_{l+1},...,A_{r}$ の中で取れる「同じ数字のペア」の最大個数を求める
$1 \le N \le 10^5$
$1 \le Q \le 10^6$
実行制限時間: 5sec.

解法

TLE解法

SegmentTreeなどでは解きづらい。
1種類の値 $x$ に限定すれば、SegmentTreeを使って区間 $[l,r]$ に含まれる $x$ の個数を取得し、 2で割って切り捨てることで $x$ で作れるペア数は $O(\log{N})$ でわかるが、それを$A_i$ の値の種類数… $O(N)$ 回繰り返さなければならない。全部で $O(QN\log{N})$ で余裕のTLE。

1つの値を1bitで管理してXORでマージすれば、「ペアにできない $x$ が存在するかどうか」がわかり、 popcountして区間長から引くことでペア数もわかるが、それもbit長が $O(N)$ になり、XORの演算に時間がかかるため無理。

公式解説の解法

Mo's Algorithm、クエリ平方分割を使う。なんか名前だけ聞いたことあった。

Mo's algorithm(クエリ平方分割)の話 - Qiita

以下のような場合に使える。

クエリを先読みできる
区間を1つ伸ばしたり縮めたりすることによる答えへの影響が簡単に計算できる

今回の場合、現在の区間内の各値の個数（の偶奇）を管理しておけば、以下の要領で区間を1つ伸縮したときの答えを更新できる。

現在の個数が奇数である値 $x$ を追加する→答えが1増える
現在の個数が偶数である値 $x$ を削除する→答えが1減る

何らかの順序でクエリを並べ替え、その順に区間を伸縮させて答えを求めていくことを考える。

例えば $r$ の昇順でクエリをソートし、この順に区間を伸縮させて答えを求めていくと、
右端の伸縮は最小限となるが、左端の伸縮で毎回 $O(N)$ かかってしまう場合がある。

1              |-|
2 |----------------|
3                 |-|
4 |-------------------|

そこで、$l$ を、なんか適当なブロックサイズ $B$ でグループ分けする。

その中ごとに $r$ の昇順でクエリを処理すると、左端の伸縮も1クエリ当たり最大 $O(B)$ で収まる。

  <-B->
1     |-|
2 |---------|
3     |--------|
4 |-----------------|

分割した分、この処理をグループの個数 $\dfrac{N}{B}$ 回繰り返す必要があるが、

$r$ を毎回左から右まで動かす計算量: $O(N \dfrac{N}{B})$
$l$ を1クエリ毎に行ったり来たりさせる計算量: $O(QB)$

全体の計算量はこの和となるので、$B=\dfrac{N}{\sqrt{Q}}$ 程度に設定すれば $O(N\sqrt{Q})$ となる。

本問題の制約では最大 $10^8$ となるが、、、まぁ $r$ を動かす範囲は右のグループほど狭まるなどの定数倍の削減があり、間に合う。