AtCoder Regular Contest++ 230

普通に数学が難しい。

A - Meeting on Tree

問題文

頂点に $1,2,\dots, N$ の番号がついた $N$ 頂点の木が与えられます.
$i=1,2,\dots, N-1$ について，$i$ 番目の辺は頂点 $u_i,v_i$ を結んでいます.
木の各頂点には $1$ 匹ずつリスがいます. リスたちは次のようにして会議を開こうとしています.
1. 会議に参加するリスを $1$ 匹以上選ぶ.
2. 選ばれたリスたちは相談し，木の頂点をひとつ選んで会議の開催地とする.
3. 選ばれたリスたちはそれぞれ，会議の開催地に到達するまで，木の辺を辿って移動する.
リスの移動には辿った辺の個数に等しいコストがかかります. 会議のコストを，選ばれたリスの移動にかかるコストの総和として定めます. リスたちは会議の開催地をうまく選ぶことで，会議のコストをできるだけ小さくしたいと考えています.
会議に参加するリスを $1$ 匹以上選ぶ方法は $2^N-1$ 通りありますが，そのそれぞれに対する「会議の開催地を適切に選んだときの会議のコストの最小値」の総和を $998244353$ で割ったあまりを求めてください.

制約

$2\le N\le 3\times 10^5$
$1\le u_i,v_i\le N$
与えられるグラフは木をなす
入力される値は全て整数

解法

主客転倒し、ある1辺に注目したとき、それが何回答えに寄与するか考える。

  X個        Y個
○○○--e--○○○
○○         ○○

ある辺を挟んで、一方に $X$ 個、もう一方に $N-X$ 個の頂点があり、そこからそれぞれ $i,j$ 個の頂点を選ぶとする。

この時、小さい方の $\min(i,j)$ 匹のリスがその辺を通って多い方に移動することになる。
$i=j$ の場合はどちらでもよいが、いずれにしろ $\min(i,j)$ 匹のリスがどっちかからどっちかへ移動する。

よって、この辺は

$\displaystyle \sum_{i=1}^{X}\sum_{j=1}^{N-X}\binom{X}{i}\binom{N-X}{j}\min(i,j)$

だけ答えに寄与することになる。

min がでてくるので上手く整理することが難しいが、小さい範囲を様々な $N,X$ で試すと当てはまる数列が見つかる。

A157933 - OEIS

OEISで説明される数列の意味は元の問題からは似ても似つかないが、どうもこの数列を三角形状に配置した $N-2$ 段目の $X-1$ 列目（0-indexed）が、まさに上記の式の値になるようだ。
$X$ は辺によって様々に変わるので、$N-2$ 段目の全ての値を前計算できれば嬉しい。

愚直に求めると $O(N^2)$ かかってしまう。
代わりに、母関数が示されているので、これを利用したい。

$\dfrac{1}{(1-2x)(1-2xy)(1-(1+y)x)}$

展開すると $1+(3+3y)x+(7+10y+7y^2)x^2+(15+25y+25y^2+15y^3)x^3+...$ となり、 $x^{N-2}$ の係数である $y$ の多項式の係数が求めたいものとなっている。
だが、やはり愚直に割り算すると $O(N^2)$ かかる。

母関数の分母は $x$ の一次式の積なので、部分分数分解する。つまり以下のような形にすることを目指す。

$\dfrac{A}{1-2x}+\dfrac{B}{1-2xy}+\dfrac{C}{1-(1+y)x}$
ヘビサイドの展開定理 | 高校数学の美しい物語

例えば $A$ は、元の母関数から $(1-2x)$ を取り去った残り $\dfrac{1}{(1-2xy)(1-(1+y)x)}$ に対して、 $1-2x=0$ となるような $x$（つまり $x=\frac{1}{2}$）を代入することで求められる。

すると、上手いこと $A,B,C$ とも共通の形が因数に現れるような形になる。それで整理すると、

$\dfrac{1}{(1-y)^2} \left ( \dfrac{2}{1-2x} + \dfrac{2y^2}{1-2xy} - \dfrac{(1+y)^2}{1-(1+y)x} \right )$

ここで、$[x^k]\frac{1}{1-ax}=a^k$ であることを利用すると、$x=N-2$ においては、

$\dfrac{2^{N-1} + 2^{N-1}y^{N} - (1+y)^{N}}{(1-y)^2}$

となる。これは、二項係数 $\binom{N}{r}$（$r=0,1,...,N$）を負にしたものに対し、両端にのみ $2^{N-1}$ を足し、それを2回、累積和をとることによって求められるとわかる。

N=5
  二項係数（負）     -1 -5 -10 -10  -5 -1
  両端に2^{N-1}加算  15 -5 -10 -10  -5 15
  累積和             15 10   0 -10 -15  0
  累積和             15 25  25  15   0      ←三角形の3段目が現れる

よって、$O(N)$ でこれを前計算し、全ての辺について寄与を足し合わせると、答えとなる。

Python3

def solve(n, edges):
    links = [[] for _ in range(n)]
    for u, v in edges:
        u -= 1
        v -= 1
        links[u].append(v)
        links[v].append(u)
    subtree_sizes = [-1] * n
    q = [0]
    while q:
        u = q[-1]
        if subtree_sizes[u] == -1:
            subtree_sizes[u] = -2
            for v in links[u]:
                if subtree_sizes[v] == -1:
                    q.append(v)
            continue
        q.pop()
        ss = 1
        for v in links[u]:
            if subtree_sizes[v] >= 0:
                ss += subtree_sizes[v]
        subtree_sizes[u] = ss

    MOD = 998244353
    facts, finvs = precompute_factorials(n + 1, MOD)

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

    table = [-ncr(n, r) for r in range(n + 1)]
    table[0] += pow(2, n - 1, MOD)
    table[n] += pow(2, n - 1, MOD)
    for i in range(n):
        table[i + 1] += table[i]
        table[i] %= MOD
    for i in range(n):
        table[i + 1] += table[i]
        table[i] %= MOD

    ans = 0
    for u, v in edges:
        su = subtree_sizes[u - 1]
        sv = subtree_sizes[v - 1]
        if su < sv:
            su, sv = sv, su
        l = n - sv
        r = sv
        k = min(l, r)
        ans += table[k - 1]
        ans %= MOD
    return ans


n = int(input())
edges = [list(map(int, input().split())) for _ in range(n - 1)]
ans = solve(n, edges)
print(ans)

目次

AtCoder Regular Contest++ 230

A - Meeting on Tree

問題文

制約

解法