ABC290 F問題 Maximum Diameter

atcoder.jp

N頂点の木の特徴

「木」は連結で閉路を持たない(無向)グラフです。

N 頂点の木は、N-1本の辺を持ちます。
1本の辺は、両端の頂点の次数にカウントされるので、 N頂点の次数を合計した値は  2(N-1)になります。
また、連結しているという条件から、次数 0 の頂点はありえず、各頂点の次数は 1以上です。
頂点が2つ以上ある木には少なくとも2個の端末点(次数1の頂点)があります。

そこで、良い木が存在するために必要な条件は、 N \ge 2 の条件の下で
 X_i \ge 1 ( i=1,2,\cdots,N)
  \sum_{i=1}^{N} X_i = 2(N-1)
となります。

ここで、直径が最大、つまり端から端までの距離が最大になる木を考えます。
 X_i ( i=1,2,\cdots,N) のうち、 X_i \ge 2 の頂点の個数を  kとします。
これらの次数2以上の頂点を一列に並べて、両端に端末点を置くと、 k+2 頂点を結ぶ道ができます。
この道の距離は、辺の数  k+1 です。

ここで、
 X_i \ge 1 ( i=1,2,\cdots,N)
 \sum_{i=1}^{N} X_i = 2(N-1)
を満たす数列  X が与えられたときに、直径 k+1の良い木が作れるか考えます。
次数2以上の頂点の個数は k、次数1の頂点の個数は  N-k です。
次数1の頂点の次数の合計は  N-k、次数2以上の頂点の次数の合計は 2(N-1)-(N-k) = N+k-2 です。
次数2以上の k個の頂点と両端の次数1の端末点で距離 k+1 の道を作ります。
道に使われなかった頂点は、次数1の頂点 N-k-2 個で、次数2以上の頂点の次数のうち、道を作るときに使われるのは
 2k であり、残り (N+k-2)-2k = N-k-2 の次数が使われずに残っています。
次数1の頂点の残りと次数3以上の頂点の次数の残りが一致しているため、次数が3以上の頂点と次数1の頂点の間に辺をつくることで余りなく木を作ることができます。

場合の数の計算

以下、 N \ge 3 とします。

 X_i \ge 1 ( i=1,2,\cdots,N)
 \sum_{i=1}^{N} X_i = 2(N-1)
を満たす整数数列  X としてありうるものすべてについて数え上げる問題になります。
このうち、 X_i \ge 2 の個数が  k 個の数列の数を数えてみます。
まず、 N個のうち、k個の X_i \ge 2N-k個の X_i = 1 に分ける場合の数を考えます。
この場合の数は、
 {}_N \mathrm{C}_k = \frac{N!}{k!(N-k)!}
です。

k個の X_i \ge 2 となる  X_i の合計は N+k-2 です。このような  X_i としてありうるものの数は、
ボールを箱に分けるときの場合の数と同じ問題とみなすことができます。最初に各箱に2個ずつボールを入れて、(残り (N+k-2)-(2k) = N-k-2)
N-k-2個のボールを k個の箱に分ける分け方の場合の数」(ボールは区別しない、箱は区別する、ボールが 0個の箱があってもよい)
という問題に帰着できます。
この答えは、
 {}_{(N-k-2)+(k)-1} \mathrm{C}_{(k)-1} = {}_{N-3} \mathrm{C}_{k-1}
です。

以上から、場合の数に、最大直径  k+1 をかけて和をとった
 \sum_{k=1}^{N-2} \left( (k+1)\cdot {}_N \mathrm{C}_k \cdot {}_{N-3} \mathrm{C}_{k-1} \right)
が答えになります。
しかし、この式のままでは時間オーバーになります。そこで、さらに式の変形を考えます。

ヴァンデルモンドの畳み込みの利用

使えそうな公式として、「ヴァンデルモンドの畳み込み」があります。
https://manabitimes.jp/math/622

和の範囲を 0 ~ n にするため、まずは  i = k-1 ,  M = N-3 と置き換えると、
 \sum_{k=1}^{N-2} \left( (k+1)\cdot {}_N \mathrm{C}_k \cdot {}_{N-3} \mathrm{C}_{k-1} \right) =
 \sum_{i=0}^{M} \left( (i+2)\cdot {}_{M+3} \mathrm{C}_{i+1} \cdot {}_{M} \mathrm{C}_{i} \right)
となります。
さらに、公式
 {}_n \mathrm{C}_r = {}_n \mathrm{C}_{n-r}
 r\cdot {}_n \mathrm{C}_r =n\cdot {}_{n-1} \mathrm{C}_{r-1}
を利用すると、
  \sum_{i=0}^{M} \left( (i+2)\cdot {}_{M+3} \mathrm{C}_{i+1} \cdot {}_{M} \mathrm{C}_{i} \right)
 = \sum_{i=0}^{M} \left( (i+1)\cdot {}_{M+3} \mathrm{C}_{i+1} \cdot {}_{M} \mathrm{C}_{i} \right) +\sum_{i=0}^{M} \left(  {}_{M+3} \mathrm{C}_{i+1} \cdot {}_{M} \mathrm{C}_{i} \right)
 = \sum_{i=0}^{M} \left( (M+3)\cdot {}_{M+2} \mathrm{C}_{i} \cdot {}_{M} \mathrm{C}_{M-i} \right) +\sum_{i=0}^{M+2} \left(  {}_{M+3} \mathrm{C}_{(M+2)-i} \cdot {}_{M} \mathrm{C}_{i} \right)
 = (M+3)\cdot {}_{2M+2} \mathrm{C}_{M} +{}_{2M+3} \mathrm{C}_{M+2}
と変形できます。 M = N-3 から、
 = N\cdot {}_{2N-4} \mathrm{C}_{N-3} +{}_{2N-3} \mathrm{C}_{N-1}
を最終的な計算式とします。

多くのテストケースで答えを求めますが、
 {}_N \mathrm{C}_k = \frac{N!}{k!(N-k)!} の値は、あらかじめ n!, 1/n! を mod 998244353 で計算しておくことで計算時間を短縮できます。

コード例 (Julia)

const MOD_ = 998244353

const NMAX = 2_000_000

fact    = Vector{Int}(undef, NMAX)
invfact = Vector{Int}(undef, NMAX)

function setup_fact( )
  fact[1] = 1
  invfact[1] = 1
  for i = 2:NMAX
    fact[i] = (i * fact[i-1]) % MOD_    # N!
    invfact[i] = invmod(fact[i],MOD_)   # 1/N!
  end
end

function Comb( N, M )
  if N < 0 || M < 0 || N - M < 0
    return 0
  end
  if M == 0 || N - M == 0
    return 1
  end

  ret = (invfact[M] * invfact[N-M])%MOD_
  ret *= fact[N]
  return ret % MOD_
end

function solve_test(N)
  ans = N * Comb(2N-4,N-3)
  ans %= MOD_
  ans += Comb(2N-3,N-1)
  ans %= MOD_

  return ans
end

function solve()

  # あらかじめ n!, 1/n! を計算
  setup_fact()

  # 入力
  readInt()  =  parse(Int,readline())

  T = readInt()

  for i = 1:T
    println(solve_test(readInt()))
  end

end  # function solve

# main
solve()