LibreOJ 2983 「WC2019」数树

计数使我快乐（

设 T1,T2 表示两棵树，E1,E2 表示它们的边集，k(E) 表示 E 中的边构成的森林的连通块个数。

对于 op=0：
这个是 PJ T2 难度吧（

对于 op=1：
首先略过 y=1 的平凡情况。
考虑对 E1∩E2 应用子集反演：

∑E2yk(E1∩E2)=∑E2yn−|E1∩E2|=∑E2∑S⊆E1∩E2∑T⊆S(−1)|S|−|T|yn−|T|=∑S⊆E1(∑T⊆S(−1)|S|−|T|yn−|T|)∑S⊆E21=∑S⊆E1yn−|S|(∑T⊆S(−y)|S|−|T|)∑S⊆E21=∑S⊆E1yn−|S|(1−y)|S|∑S⊆E21

设 S 中的边构成 k=n−|S| 个连通块，其点数分别为 a1,a2,…,ak。
那么有结论 ∑S⊆E21=nk−2k∏i=1ai

证明可以参见 rqy 的博客或者 shadowice1984 的题解，这里不再赘述了。
（Prufer 序列好！）

于是 ∑S⊆E1yn−|S|(1−y)|S|∑S⊆E21=∑S⊆E1yk(1−y)n−knk−2k∏i=1ai=(1−y)nn2∑S⊆E1k∏i=1ny1−yai

设 c=ny1−y，考虑后一部分。
相当于在 T1 中划分出 k 个连通块，第 i 个连通块有 cai 的贡献。

这一问题有显然的 O(n2) 的树形 DP 解法。
而可以像 rqy 一样使用生成函数进行优化 DP 做到 O(n)。
但这里给出一个更为简单的做法。

考虑 cai 的组合意义，相当于在每个连通块中选一个点，每个点有 c 的贡献。
于是可以简单地设 fu,0/1 表示 u 所在的连通块是否选出了这个点的贡献和。
同样是 O(n) 的。

对于 op=2：
类比上一部分，此处略去推导过程，则需要计算 (1−y)nn4∑Sk∏i=1n2y1−ya2i

对于后一部分的和式，考虑枚举所有的 k 和 a1,a2,…,ak。
n∑k=1∑a1+a2+⋯+ak=nn!k!∏ki=1ai!k∏i=1n2y1−ya2i⋅aai−2i=n!n∑k=11k!∑a1+a2+⋯+ak=nk∏i=1n2y(1−y)ai!aaii

考虑构造生成函数来计算这个东西，设 F(x)=n2y1−y∞∑i=1iii!xi

则答案等于 (1−y)nn!n4[xn]n∑k=11k!Fk(x)

注意到这是一个 exp 的形式，所以答案就是 (1−y)nn!n4[xn]expF(x)

代码：

#include <cstdio>
#include <vector>
#include <cstring>
#include <algorithm>
#include <unordered_map>
#include <unordered_set>
#define add(a,b) (a + b >= mod ? a + b - mod : a + b)
#define dec(a,b) (a < b ? a - b + mod : a - b)
using namespace std;
const int N = 1e5;
const int mod = 998244353;
inline int fpow(int a,int b)
{
    int ret = 1;
    for(;b;b >>= 1)
        (b & 1) && (ret = (long long)ret * a % mod),a = (long long)a * a % mod;
    return ret;
}
int n,y,op;
namespace Poly
{
    const int N = 1 << 18;
    const int G = 3;
    int lg2[N + 5];
    int rev[N + 5],fac[N + 5],ifac[N + 5],inv[N + 5];
    int rt[N + 5],irt[N + 5];
    inline void init()
    {
        for(register int i = 2;i <= N;++i)
            lg2[i] = lg2[i >> 1] + 1;
        int w = fpow(G,(mod - 1) / N);
        rt[N >> 1] = 1;
        for(register int i = (N >> 1) + 1;i <= N;++i)
            rt[i] = (long long)rt[i - 1] * w % mod;
        for(register int i = (N >> 1) - 1;i;--i)
            rt[i] = rt[i << 1];
        fac[0] = 1;
        for(register int i = 1;i <= N;++i)
            fac[i] = (long long)fac[i - 1] * i % mod;
        ifac[N] = fpow(fac[N],mod - 2);
        for(register int i = N;i;--i)
            ifac[i - 1] = (long long)ifac[i] * i % mod;
        for(register int i = 1;i <= N;++i)
            inv[i] = (long long)ifac[i] * fac[i - 1] % mod;
    }
    struct poly
    {
        vector<int> a;
        inline poly(int x = 0)
        {
            x && (a.push_back(x),1);
        }
        inline poly(const vector<int> &o)
        {
            a = o,shrink();
        }
        inline void shrink()
        {
            for(;!a.empty() && !a.back();a.pop_back());
        }
        inline int size() const
        {
            return a.size();
        }
        inline void resize(int x)
        {
            a.resize(x);
        }
        inline int operator[](int x) const
        {
            if(x < 0 || x >= size())
                return 0;
            return a[x];
        }
        inline int &operator[](int x)
        {
            return a[x];
        }
        inline void clear()
        {
            vector<int>().swap(a);
        }
        inline void ntt(int type = 1)
        {
            int n = size();
            type == -1 && (reverse(a.begin() + 1,a.end()),1);
            int lg = lg2[n] - 1;
            for(register int i = 0;i < n;++i)
                rev[i] = (rev[i >> 1] >> 1) | ((i & 1) << lg),
                i < rev[i] && (swap(a[i],a[rev[i]]),1);
            for(register int w = 2,m = 1;w <= n;w <<= 1,m <<= 1)
                for(register int i = 0;i < n;i += w)
                    for(register int j = 0;j < m;++j)
                    {
                        int t = (long long)rt[m | j] * a[i | j | m] % mod;
                        a[i | j | m] = dec(a[i | j],t),a[i | j] = add(a[i | j],t);
                    }
            if(type == -1)
                for(register int i = 0;i < n;++i)
                    a[i] = (long long)a[i] * inv[n] % mod;
        }
        friend inline poly operator+(const poly &a,const poly &b)
        {
            vector<int> ret(max(a.size(),b.size()));
            for(register int i = 0;i < ret.size();++i)
                ret[i] = add(a[i],b[i]);
            return poly(ret);
        }
        friend inline poly operator-(const poly &a,const poly &b)
        {
            vector<int> ret(max(a.size(),b.size()));
            for(register int i = 0;i < ret.size();++i)
                ret[i] = dec(a[i],b[i]);
            return poly(ret);
        }
        friend inline poly operator*(poly a,poly b)
        {
            if(a.a.empty() || b.a.empty())
                return poly();
            int lim = 1,tot = a.size() + b.size() - 1;
            for(;lim < tot;lim <<= 1);
            a.resize(lim),b.resize(lim);
            a.ntt(),b.ntt();
            for(register int i = 0;i < lim;++i)
                a[i] = (long long)a[i] * b[i] % mod;
            a.ntt(-1),a.shrink();
            return a;
        }
        poly &operator+=(const poly &o)
        {
            resize(max(size(),o.size()));
            for(register int i = 0;i < o.size();++i)
                a[i] = add(a[i],o[i]);
            return *this;
        }
        poly &operator-=(const poly &o)
        {
            resize(max(size(),o.size()));
            for(register int i = 0;i < o.size();++i)
                a[i] = dec(a[i],o[i]);
            return *this;
        }
        poly &operator*=(poly o)
        {
            return (*this) = (*this) * o;
        }
        poly deriv() const
        {
            if(a.empty())
                return poly();
            vector<int> ret(size() - 1);
            for(register int i = 0;i < size() - 1;++i)
                ret[i] = (long long)(i + 1) * a[i + 1] % mod;
            return poly(ret);
        }
        poly integ() const
        {
            if(a.empty())
                return poly();
            vector<int> ret(size() + 1);
            for(register int i = 0;i < size();++i)
                ret[i + 1] = (long long)a[i] * inv[i + 1] % mod;
            return poly(ret);
        }
        inline poly modxn(int n) const
        {
            n = min(n,size());
            return poly(vector<int>(a.begin(),a.begin() + n));
        }
        inline poly inver(int m) const
        {
            poly ret(fpow(a[0],mod - 2));
            for(register int k = 1;k < m;)
                k <<= 1,ret = (ret * (2 - modxn(k) * ret)).modxn(k);
            return ret.modxn(m);
        }
        inline poly log(int m) const
        {
            return (deriv() * inver(m)).integ().modxn(m);
        }
        inline poly exp(int m) const
        {
            poly ret(1);
            for(register int k = 1;k < m;)
                k <<= 1,ret = (ret * (1 - ret.log(k) + modxn(k))).modxn(k);
            return ret.modxn(m);
        }
    };
}
using Poly::init;
using Poly::poly;
namespace Task0
{
    unordered_map<int,unordered_set<int>> vis;
    int ans,k;
    void main()
    {
        k = n;
        int u,v;
        for(register int i = 2;i <= n;++i)
            scanf("%d%d",&u,&v),u > v && (swap(u,v),1),vis[u].insert(v);
        for(register int i = 2;i <= n;++i)
            scanf("%d%d",&u,&v),u > v && (swap(u,v),1),k -= vis.count(u) && vis[u].count(v);
        ans = fpow(y,k);
        printf("%d\n",ans);
    }
}
namespace Task1
{
    int to[(N << 1) + 5],pre[(N << 1) + 5],first[N + 5];
    inline void add_edge(int u,int v)
    {
        static int tot = 0;
        to[++tot] = v,pre[tot] = first[u],first[u] = tot;
    }
    int fa[N + 5];
    int c,ans,f[N + 5][2];
    void dfs(int p)
    {
        f[p][0] = 1,f[p][1] = c;
        for(register int i = first[p];i;i = pre[i])
            if(to[i] ^ fa[p])
                fa[to[i]] = p,
                dfs(to[i]),
                f[p][1] = ((long long)f[p][0] * f[to[i]][1] + (long long)f[p][1] * f[to[i]][0] + (long long)f[p][1] * f[to[i]][1]) % mod,
                f[p][0] = ((long long)f[p][0] * f[to[i]][0] + (long long)f[p][0] * f[to[i]][1]) % mod;
    }
    void main()
    {
        if(y == 1)
        {
            printf("%d\n",fpow(n,n - 2));
            return ;
        }
        int u,v;
        for(register int i = 2;i <= n;++i)
            scanf("%d%d",&u,&v),add_edge(u,v),add_edge(v,u);
        c = (long long)n * y % mod * fpow(dec(1,y),mod - 2) % mod;
        dfs(1);
        ans = (long long)fpow(dec(1,y),n) * fpow(n,mod - 3) % mod * f[1][1] % mod;
        printf("%d\n",ans);
    }
}
namespace Task2
{
    poly f;
    int c,ans;
    void main()
    {
        if(y == 1)
        {
            printf("%d\n",fpow(n,2 * n - 4));
            return ;
        }
        init();
        c = (long long)n * n % mod * y % mod * fpow(dec(1,y),mod - 2) % mod;
        f.resize(n + 1);
        for(register int i = 1;i <= n;++i)
            f[i] = (long long)c * fpow(i,i) % mod * Poly::ifac[i] % mod;
        f = f.exp(n + 1),ans = (long long)fpow(dec(1,y),n) * Poly::fac[n] % mod * fpow(n,mod - 5) % mod * f[n] % mod;
        printf("%d\n",ans);
    }
}
int main()
{
    scanf("%d%d%d",&n,&y,&op);
    if(!op)
        Task0::main();
    else if(op == 1)
        Task1::main();
    else
        Task2::main();
}