AtCoder Beginner Contest 409 F,G問題メモ

AtCoder Beginner Contest 409

F - Connecting Points

問題文

$2$ 次元平面上に $N$ 頂点 $0$ 辺のグラフ $G$ があります。頂点には $1$ から $N$ までの番号が付いており、頂点 $i$ は座標 $(x_i,y_i)$ にあります。
$G$ の頂点 $u,v$ に対して、$u,v$ の距離 $d(u,v)$ をマンハッタン距離 $d(u,v)=|x_u-x_v|+|y_u-y_v|$ で定義します。
また、$G$ の $2$ つの連結成分 $A,B$ に対して、$A,B$ の頂点集合をそれぞれ $V(A),V(B)$ としたとき、$A,B$ の距離 $d(A,B)$ を $d(A,B)=\min\lbrace d(u,v)\mid u \in V(A), v\in V(B)\rbrace$ で定義します。
以下で説明されるクエリを $Q$ 個処理してください。クエリは次の $3$ 種類のいずれかです。
- '1 a b' : $G$ の頂点数を $n$ としたとき、頂点 $n+1$ の座標を $(x_{n+1},y_{n+1})=(a,b)$ として、$G$ に頂点 $n+1$ を追加する。
- '2' : $G$ の頂点数を $n$、連結成分数を $m$ とする。
  - $m=1$ のとき、'-1' を出力する。
  - $m\geq 2$ のとき、距離が最小である連結成分をすべてマージし、その最小の距離の値を出力する。
- '3 u v' : 頂点 $u,v$ が同じ連結成分にあるならば 'Yes' を、そうでないならば 'No' を出力する。

制約

$2\leq N\leq 1500$
$1\leq Q\leq 1500$
$0\leq x_i,y_i\leq 10^9$
$1$ 種類目のクエリについて、$0\leq a,b\leq 10^9$
$3$ 種類目のクエリについて、そのクエリを処理する直前の $G$ の頂点数を $n$ としたとき、$1\leq u\lt v\leq n$
入力は全て整数

解法

一見複雑だが、解法は素直。
Heap Queue に、条件を満たすものも満たさないものもとりあえず突っ込んでおいて、取り出すときにチェックする（条件を満たすものが見つかるまで pop し続ける）という操作は他の問題でもよく出てきがち。

最終的な頂点数の上限は $O(N+Q)$ であり、全頂点間の距離 $O((N+Q)^2)$ を求めても十分間に合う。

初期状態で $G$ にある $N$ 頂点の全点間距離をHeap Queue $H$ につっこんでおく。
また、連結状態を Union-Find $T$ で管理する。

クエリ1では、新たに追加された頂点と、今ある $n$ 頂点の距離を $H$ に追加する。

クエリ2では、まだ連結でない頂点対 $(i,j)$ が見つかるまで、$H$ で距離の小さい方からpopし続ける。
見つかったらそれを $k$ とし、$T$ で $(i,j)$ を結ぶ。
その後、$H$ に含まれる最小距離が $k$ である間、追加でpopし続け、その頂点も全部結ぶ。
$H$ が空になるまで見つからなかったら、全て連結ということなので、$-1$ を出力して終了。

クエリ3は、普通に $T$ 上で $(u,v)$ が連結かを調べる。

以上で $O((N+Q)^2 \log(N+Q))$ でできる。

Python3

from heapq import heappush, heappop

from atcoder.dsu import DSU


class KeyHeapQueue:
    def __init__(self):
        self.hq = []
        self.mask = (1 << 12) - 1

    def heappush(self, d, i, j):
        key = (d << 24) | (i << 12) | j
        heappush(self.hq, key)

    def heappop(self):
        key = heappop(self.hq)
        j = key & self.mask
        i = (key >> 12) & self.mask
        d = key >> 24
        return d, i, j

    def get_top(self):
        if len(self.hq) == 0:
            return -1, -1, -1
        key = self.hq[0]
        j = key & self.mask
        i = (key >> 12) & self.mask
        d = key >> 24
        return d, i, j


def solve(n, q, points, queries):
    dsu = DSU(n + q)
    hq = KeyHeapQueue()
    next_i = n
    for i in range(n):
        for j in range(i + 1, n):
            x1, y1 = points[i]
            x2, y2 = points[j]
            hq.heappush(abs(x1 - x2) + abs(y1 - y2), i, j)

    ans = []
    for query in queries:
        if query[0] == 1:
            j = next_i
            next_i += 1
            x2, y2 = query[1], query[2]
            for i, (x1, y1) in enumerate(points):
                hq.heappush(abs(x1 - x2) + abs(y1 - y2), i, j)
            points.append((x2, y2))
        elif query[0] == 2:
            k = -1
            while len(hq.hq) > 0:
                d, i, j = hq.heappop()
                if dsu.same(i, j):
                    continue
                k = d
                dsu.merge(i, j)
                break
            if k == -1:
                ans.append('-1')
                continue
            ans.append(str(k))
            while len(hq.hq) > 0:
                d, i, j = hq.get_top()
                if d > k:
                    break
                dsu.merge(i, j)
                hq.heappop()
        else:
            u = query[1] - 1
            v = query[2] - 1
            if dsu.same(u, v):
                ans.append('Yes')
            else:
                ans.append('No')

    return ans


n, q = map(int, input().split())
points = [tuple(map(int, input().split())) for _ in range(n)]
queries = [tuple(map(int, input().split())) for _ in range(q)]
ans = solve(n, q, points, queries)
print('\n'.join(ans))

G - Accumulation of Wealth

問題文

$2$ 以上の整数 $N$ と、$0$ 以上 $100$ 以下の整数 $P$ が与えられます。以下、 $p=P/100$ とします。
数列 $A$ があります。はじめ、$A$ の長さは $1$ で、唯一の要素は $1$ です。
数列 $A$ に対して、以下の操作を $N-1$ 回繰り返します。
- $A$ に現れない最小の正整数を $m$ とする。確率 $p$ で操作 $1$、確率 $1-p$ で操作 $2$ を行う：
  - 操作 $1$: $A$ の末尾に $m$ を追加する。
  - 操作 $2$: $1,2,\dots,m-1$ が $A$ に含まれる個数を $c_1,c_2,\dots,c_{m-1}$ とする。$1$ 以上 $m-1$ 以下の整数 $k$ を $c_k$ に比例する確率で $1$ つ選ぶ。すなわち、確率 $c_k/\sum_{j=1}^{m-1} c_j$ で $k$ を選ぶ。そして、$A$ の末尾に $k$ を追加する。

各 $k=1,2,\dots,N$ について、$N-1$ 回の操作後に $A$ に含まれる $k$ の個数の期待値を $\mathrm{mod}\; 998244353$ で求めてください。

制約

$2 \leq N \leq 10^5$
$0 \leq P \leq 100$
入力はすべて整数

解法

数式をこねくり回して畳み込みの形に持っていく系問題。
まぁそうだろうなとは思いつつ、どうやったら畳み込みの形になるか分からなかった。

立式

$p=\frac{P}{100}$、$q=1-p$ とする。

$E(i,k)$ を「$i$ 回の操作後の $k$ の個数期待値」とする。$E(N-1,k)$ を各 $k$ につき求めたい。

ひとまず最初から存在する $k=1$ は別として、$k \ge 2$ に対して以下のアプローチを取る。

「① $j$ 回目にはじめて $k$ が追加される確率」と「② 追加された1個から、残り $N-j-1$ 回の操作で期待値が何倍になるか」を分けて考える。

①は、「$j-1$ 回目までに操作1が $k-2$ 回選ばれた後、$j$ 回目で操作1が選ばれる」という事象なので、以下になる。

${}_{j-1}\mathrm{C}_{k-2} p^{k-2}q^{j-k+1} \times p$
ただし、$j-1 \lt k-2$ の場合は $0$

②は、$k$ に依らない。$k=2$ であろうと $3$ であろうと、 $j$ 回目に追加された1個が最後までに期待値何個になるかは同じである。

①の「$j$ 回目に始めて $k$ が追加」が発生したという条件の下の $E(i,k)$ を $E'_{j}(i)$ とする。
$k$ は何であっても同じなので省略する。以下が成り立つ。

$E' = E \times (①の式)$
$E'_{j}(j)=1$

$i+1$ 回目に追加されるのが $k$ である確率は、まず確率 $q$ で操作2が選ばれ、次に $i$ 個の要素から $E'_{j}(i)$ 個ある $k$ が選ばれなければならないので、$q \times \frac{E'_{j}(i)}{i}$ である。

操作1が選ばれたり、操作2で他の要素が選ばれたときは $E'_{j}(i)$ のままなので、あわせると

$E'_{j}(i+1) = E'_{j}(i) \times (1+\frac{q}{i})$

となる。この遷移を $N-1$ 回目まで繰り返すと、結局

$\displaystyle E'_{j}(N-1) = \prod_{l=j}^{N-2} (1+\frac{q}{l})$

となる。$E'_{N-1}(N-1)=1$ よりはじめて、$j=N-2,...,1$ と遡るたびに $1+\frac{q}{j}$ を乗算していくことで、 $E'_{j}(N-1)$ を各 $j$ に対して $O(N)$ で求めることができる。

①と②をあわせ、初登場までにかかった操作回数 $j$ によって場合分けして、$E(N-1,k)$ は以下のようになる。

$\displaystyle E(N-1,k) = \sum_{j=1}^{N-1} {}_{j-1}\mathrm{C}_{k-2} p^{k-2}q^{j-k+1} \times p \times E'_{j}(N-1)$

これを各 $k$ に対して求めると答えだが、このままでは $O(N^2)$ となる。まだ高速化の必要がある。

形式的冪級数で考えると、二項係数を数式の累乗にまとめられるので、上の式は

$\displaystyle F(x) = \sum_{j=1}^{N-1} px(px+q)^{j-1} \times \prod_{l=j}^{N-2} (1+\frac{q}{l})$
$k=2,3,...,N$ の時の答えは、$F(x)$ の $x^1,x^2,...,x^{N-1}$ の係数

となる。

Polynomial Taylor Shift

ここで、Polynomial Taylor Shift というテクニックが使える。

これは、$x$ の $N$ 次多項式 $A(x)=a_0+a_1x+a_2x^2+...+a_Nx^N$ と定数 $c$ が与えられたとき、 $A(x+c) = b_0+b_1x+b_2x^2+...+b_Nx^N$ と表される $(b_0,b_1,...,b_N)$ を $O(N \log{N})$ で求めるテクニックである。

Polynomial Taylor Shift の式変形自体はけんちょん氏の記事に詳しいのでここでは割愛。

ともかく、これができるとして、

$\displaystyle A(x)=\sum_{j=1}^{N-1} p^j \prod_{l=j}^{N-2} (1+\frac{q}{l}) x^j$
$c=\frac{q}{p}$

とすると、上記の $F(x)$ にも Polynomial Taylor Shift を適用できる。
この式だと、$x$ の係数が1つずれて、$x^0,x^1,...,x^{N-2}$ の係数が答えとなる。

最後に、$k=1$ の時の答えは、全体の答えの総和が $N$ になることを利用してもいいし、また $E'_1(N-1)$ であることを利用して計算してもよい。

Python3

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]


def polynomial_taylor_shift(
        aaa: npt.NDArray[np.int64],
        c: int,
        MOD: int,
        facts: npt.NDArray[np.int64],
        finvs: npt.NDArray[np.int64],
) -> npt.NDArray[np.int64]:
    """
    A = (a0,a1,...,aN) と 定数 c が与えられたとき、
    f(x)   = a0 + a1*x + a2*x^2 + ... + aN*x^N に対して、
    f(x+c) = b0 + b1*x + b2*x^2 + ... + bN*x^N の各係数
    B = (b0,b1,...,bN) を O(NlogN) で算出する。

    畳み込みを行う関数 convolve(fff, ggg) が定義されていること。
    （関数名が違う場合は適宜置換して）

    See Also:
        https://drken1215.hatenablog.com/entry/2023/09/08/003100
    """
    n = aaa.size - 1
    fff = aaa * facts[:n + 1] % MOD
    ccc = [1]
    for _ in range(n):
        ccc.append(ccc[-1] * c % MOD)
    ccc = np.array(ccc, dtype=np.int64)
    ggg = (ccc * finvs[:n + 1])[::-1] % MOD
    hhh = convolve_998244353(fff, ggg)[n:]
    result = hhh * finvs[:n + 1] % MOD
    return result


def precompute_factorials(n, MOD):
    f = 1
    factorials = [1]
    for m in range(1, n + 1):
        f = f * m % MOD
        factorials.append(f)
    f = pow(f, MOD - 2, MOD)
    finvs = [1] * (n + 1)
    finvs[n] = f
    for m in range(n, 1, -1):
        f = f * m % MOD
        finvs[m - 1] = f
    return factorials, finvs


def solve(n, p):
    if p == 0:
        ans = [0] * n
        ans[0] = n
        return ans

    MOD = 998244353
    p *= pow(100, -1, MOD)
    p %= MOD
    q = (1 - p) % MOD

    facts, finvs = precompute_factorials(n, MOD)
    facts_np = np.array(facts, dtype=np.int64)
    finvs_np = np.array(finvs, dtype=np.int64)

    prod = [1]
    for i in range(n - 2):
        prod.append(prod[-1] * (1 + q * finvs[n - 1 - i] * facts[n - 2 - i] % MOD) % MOD)
    prod.reverse()

    aaa = []
    pow_p = p
    for i in range(n - 1):
        aaa.append(pow_p * prod[i] % MOD)
        pow_p *= p
        pow_p %= MOD
    aaa = np.array(aaa, dtype=np.int64)
    c = q * pow(p, -1, MOD) % MOD

    bbb = polynomial_taylor_shift(aaa, c, MOD, facts_np, finvs_np)
    b0 = prod[0] * (1 + q) % MOD
    ans = [b0] + bbb.tolist()
    return ans


n, p = map(int, input().split())
ans = solve(n, p)
print('\n'.join(map(str, ans)))

目次

AtCoder Beginner Contest 409 F,G問題メモ

F - Connecting Points

問題文

制約

解法

G - Accumulation of Wealth

問題文

制約

解法

立式

Polynomial Taylor Shift