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

AtCoder Beginner Contest 436

ライブラリ準備不足！

E - Minimum Swap

問題文

$(1,2,\ldots,N)$ を並べ替えた整数列 $P=(P _ 1,P _ 2,\ldots,P _ N)$ が与えられます。
あなたは、次の操作を $0$ 回以上行って $P$ を列 $(1,2,\ldots,N)$ と一致させたいです。
- $1\le i\lt j\le N$ を満たす整数の組 $(i,j)$ をひとつ選ぶ。$P _ i$ と $P _ j$ の値を入れ替える。

$P$ を列 $(1,2,\ldots,N)$ と一致させるために必要な最小の操作回数を $K$ 回とします。
$K$ 回の操作で $P$ を列 $(1,2,\ldots,N)$ と一致させるような操作列の $\mathbf{1}$ 回目の操作としてあり得る操作の数を求めてください。
ただし、$2$ つの操作は選んだ整数の組 $(i,j)$ が異なるとき、またそのときに限り区別します。

制約

$2\le N\le3\times10 ^ 5$
$1\le P _ i\le N\ (1\le i\le N)$
$P _ i\ne P _ j\ (1\le i\lt j\le N)$
$i\ne P _ i$ を満たす $1\le i\le N$ が存在する
入力はすべて整数

解法

順列上で任意swapを繰り返してソートする問題。ここ数年のABCでは頻出のテーマ。

頂点 $1～N$ を用意し、$i→P_i$ に辺を張ったグラフを考えると、いくつかのサイクルができる。

同じサイクル上の任意の2点 $(i,j)$ を $P$ で入れ替えると、対応するグラフではサイクルが必ず1個増える。
最終的にサイクルが $N$ 個になったら完成。

よって、サイクル毎に独立に考えられる。
$c$ 頂点のサイクルからは、$\frac{c(c-1)}{2}$ 通りの $(i,j)$ の選び方がある。
サイクル毎にサイズ $c$ を調べて、上の計算結果を合計したら答え。

Python3

from atcoder.dsu import DSU


def solve(n, ppp):
    uft = DSU(n)
    for i, p in enumerate(ppp):
        p -= 1
        uft.merge(i, p)

    ans = 0
    for grp in uft.groups():
        s = len(grp)
        ans += s * (s - 1) // 2

    return ans


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

F - Starry Landscape Photo

問題文

AtCoder 星から見える夜空には $N$ 個の星があり、これらの $N$ 個の星は東から西へ一直線上に並んでいます。
東から $i$ 番目 $(1\le i\le N)$ の星は、これらの星の中で $B _ i$ 番目に明るい星です。
高橋くんは、次のような手順で夜空の写真を撮ることにしました。
- $1\le l\le r\le N$ を満たす整数の組 $(l,r)$ を選び、東から $l$ 番目、$l+1$ 番目、$\ldots$、$r$ 番目の星が全てフレームに収まり、他の星がフレームに入らないようにカメラを設置する。
- $1\le b\le N$ を満たす整数 $b$ を選び、$N$ 個の星のうち明るさが $1$ 番目から $b$ 番目に入る（かつフレームに収まっている）星がすべて写り、そうでない星が写らないようにシャッターを開放する。
ただし、星が $1$ つも写らないように写真を撮ることはできません。
このようにして撮影された夜空の写真に写っている星の集合としてありえるものが何通りあるか求めてください。

制約

$1\le N\le5\times10 ^ 5$
$1\le B _ i\le N\ (1\le i\le N)$
$B _ i\ne B _ j\ (1\le i\lt j\le N)$
入力はすべて整数

解法

カメラに「光量を絞る」機能があり、「フレームに入っていても映らない星があるように自由に明るさを調節できる」ことを理解していないと、問題文の意図が若干伝わりづらいか。（最近のデジカメは気にしなくても最適な明るさにしてくれ、意識して光量を調節することが稀なので）
まぁ、サンプルを見ればすぐ分かるようにしてくれている。

カメラに写っている最も暗い星で場合分けする。

この時、写真に写っている (最も左の星、最も右の星) の組の個数だけ、異なる星の集合ができる。

$L_i:=$ 星 $i$ より左で、$B_i$ より明るい星の個数
$R_i:=$ 星 $i$ より右で、$B_i$ より明るい星の個数

とすると、最も暗い星を $i$ としたときの答えへの寄与は

$(L_i+1)\times(R_i+1)$

で求められる。

転倒数を求める時のように FenwickTree などを使えば、$L_i,R_i$ を $O(N \log{N})$ で計算できる。

Python3

from atcoder.fenwicktree import FenwickTree


def solve(n, bbb):
    fwt1 = FenwickTree(n)
    fwt2 = FenwickTree(n)
    bbb = [b - 1 for b in bbb]
    left = [0] * n
    right = [0] * n
    for i in range(n):
        b = bbb[i]
        left[i] = fwt1.sum(0, b)
        fwt1.add(b, 1)
    for i in range(n - 1, -1, -1):
        b = bbb[i]
        right[i] = fwt2.sum(0, b)
        fwt2.add(b, 1)
    ans = 0
    for i in range(n):
        ans += (left[i] + 1) * (right[i] + 1)
    return ans


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

G - Linear Inequation

問題文

長さ $N$ の正整数列 $A=(A _ 1,A _ 2,\ldots,A _ N)$ および正整数 $M$ が与えられます。
次の条件を満たす長さ $N$ の非負整数列 $x=(x _ 1,x _ 2,\ldots,x _ N)$ がいくつあるか求めてください。
- $\displaystyle\sum _ {i=1} ^ NA _ ix _ i\le M$

答えは非常に大きくなる場合があるので、答えを $998244353$ で割った余りを求めてください。

制約

$1\le N\le100$
$1\le A _ i\le100\ (1\le i\le N)$
$1\le M\le10 ^ {18}$
入力はすべて整数

解法

最近、FPS 24 題という、形式的冪級数をテーマとしたコンテストがあった。
これのおかげで、解法を連想しやすかった人が増えたかもしれない。

問題文中の $x$ と被るのでややこしいが、今後、$x$ は形式的冪級数に用いる変数を指すとする。

形式的冪級数に置き換えて考えると、

$F_1(x)=1+x^{A_1}+x^{2A_1}+x^{3A_1}+...$
$F_2(x)=1+x^{A_2}+x^{2A_2}+x^{3A_2}+...$
：
$F_N(x)=1+x^{A_N}+x^{2A_N}+x^{3A_N}+...$

という多項式 $N$ 個の積 $F_1F_2...F_N(x)$ の、$x^0～x^M$ の係数の合計が答えとなる。

このような指数が等間隔に無限に続く式は、$F_i(x)=\dfrac{1}{1-x^{A_i}}$ というようにも表すことができる。
また、数列の初項からの累積和は、形式的冪級数としては $\dfrac{1}{1-x}$ をかけることに相当する。
よって求めたい多項式は

$G(x)=\dfrac{1}{(1-x)(1-x^{A_1})(1-x^{A_2})...(1-x^{A_N})}$

のように表せる。

$G(x)$ の分母の次数は制約より、高々 $O(N \max(A)) \simeq 10^4$ である。
1つずつ畳み込みでマージしていけば、$O(N^2 \max(A) \log{(N \max(A))})$ で、全て展開した形を求められる。

$G(x)$ の $x^M$ の項の係数を求めたい。

母関数が分かっている数列の $M$ 項目を高速に求める方法として、Bostan-Mori がある。
母関数の次数を $d$ として、$O(d \log{d} \log{M})$ で求められる。

Python3

Numbaでは数列 $Q$ をコピーするとき、Q_minus = Q[:] としてしまうと、実際はコピーされない（Q_minus への変更が Q にも反映されてしまう）という罠がある。前もハマったような気がしたんだけど、ついやってしまう。

from collections import deque

import numpy as np
import numpy.typing as npt
import numba


@numba.njit('(i8[:],i8[:])', cache=True)
def convolve_998244353(aaa: npt.NDArray[np.int64], bbb: npt.NDArray[np.int64]) -> npt.NDArray[np.int64]:
    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 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]


@numba.njit('(i8,i8[:],i8[:])', cache=True)
def bostan_mori(N, P, Q):
    """
    Bostan-Mori 法を用いて [x^N] P(x)/Q(x) を計算する。（mod 998244353 固定）

    Args:
        N (int): 求めたい項のインデックス (0-indexed)
        P (list): 分子多項式の係数リスト [p0, p1, ...]
        Q (list): 分母多項式の係数リスト [q0, q1, ...]
    Returns:
        int: 数列の第 N 項
    """
    MOD = 998244353

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

    # N が 0 になるまでループ
    while N > 0:
        # Q(-x) を生成
        # 奇数次の係数の符号を反転させる (mod を考慮)
        Q_minus = Q.copy()
        for i in range(1, len(Q_minus), 2):
            Q_minus[i] = (MOD - Q_minus[i]) % MOD

        # 畳み込み (多項式乗算)
        # U(x) = P(x) * Q(-x)
        U = convolve_998244353(P, Q_minus)
        # V(x) = Q(x) * Q(-x) = V(x^2) となるはずの多項式
        V = convolve_998244353(Q, Q_minus)

        # N の偶奇に応じて P を更新し、次数を下げる
        if N % 2 == 0:
            # N が偶数: U の偶数次の係数を取り出す (x -> x^2 の置換に対応)
            P = U[0::2]
        else:
            # N が奇数: U の奇数次の係数を取り出す
            P = U[1::2]

        # Q を更新 (V の偶数次の係数を取り出す)
        Q = V[0::2]

        # P が空になったら、それ以降の係数は 0
        if len(P) == 0:
            return 0

        # N を半分にする
        N >>= 1

    # ベースケース N=0: P(0) / Q(0)
    # 答えは P[0] * Q[0]^(-1) (mod p)
    inv_Q0 = mod_pow(Q[0], MOD - 2)
    return (P[0] * inv_Q0) % MOD


def solve(n, m, aaa):
    q = deque()
    q.append(np.array([1, -1], np.int64))
    for a in aaa:
        arr = np.zeros(a + 1, np.int64)
        arr[0] = 1
        arr[a] = -1
        q.append(arr)

    while len(q) > 1:
        a1 = q.popleft()
        a2 = q.popleft()
        a3 = convolve_998244353(a1, a2)
        q.append(a3)

    q = q[0]
    p = np.ones(1, np.int64)
    p[0] = 1
    ans = bostan_mori(m, p, q)
    return ans


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

目次

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

E - Minimum Swap

問題文

制約

解法

F - Starry Landscape Photo

問題文

制約

解法

G - Linear Inequation

問題文

制約

解法