PKUWC2019 D1T2

### 题目描述
给你一棵n个点的树，每一个点有个颜色$c_i$。定义一种颜色的树是：把该颜色的点两两路径上的点和边都拿出来构成的图形。你可以选出$k$种颜色来，求：有多少种方案可以使每种颜色交集产生的树非空。对$k\in [1,m]$分别回答。

### 输入格式
第一行两个整数$n,m$
第二行$n$个数表示每个点的颜色
第三行$n-1$个数，第$i$个表示$i+1$的父亲

### 输出格式
一行$m$个数，表示答案

### 数据范围及子任务
$m\le n\le 100000$
$Sub1:n,m\le 100(12pts)$
$Sub2:m<=2(41pts)$
$Sub3:m<=5000(14pts)$
$Sub4:无特殊限制$

题解：首先转化一下，考虑一个方案的答案可以转化为$交集点数-交集边数$
现在我们对点数/边数分别求解
以点为例子，考虑一个点在多少方案中被计算
这个很简单，只需要找出一个点在几种颜色的树里面，剩下的算一个组合数即可。找出一个点在几棵树里面可以建虚树然后树链剖分加差分。

接下来设在$i$种颜色的树中的点有$t_i$个，发现答案是：
$$ans_k=\sum^{n}_{i=k}t_i\binom i k$$
$$ans_k=\frac{1}{k!}\sum^{n}_{i=k}\frac{t_i\times i!}{(i-k)!}$$
$$ans_k=\frac{1}{k!}\sum^{n-k}_{i=0}\frac{t_{i+k}\times (i+k)!}{i!}$$
发现类似通配符匹配，可以构造卷积：
$$f_i=\frac{1}{(n-i)!},g_i=t_i\times i!$$
答案就是$g\times f$
对边做同样的事情，两式相减即可。

```
#include<cstdio>
#include<vector>
#include<algorithm>
#include<cstring>
using namespace std;
const int mod = 998244353;
const int maxn = 1e5 + 5;
const int G = 3;
vector<int > nds[maxn];
void Add(int &x, const int &a) {
    x += a; if(x >= mod) x -= mod;
}
void Dec(int &x, const int &a) {
    x -= a; if(x < 0) x += mod;
}
int mul(const int &a, const int &b) {
    return 1ll * a * b% mod;
}
int fpow(int a, int b) {
    int ans = 1;
    while(b) {
        if(b & 1) ans = mul(ans, a);
        a = mul(a, a), b >>= 1;
    }
    return ans;
}
namespace Poly {
    const int maxd = 1e6 + 5;
    struct PLG {
        vector<int > x;
        int deg() const {return x.size() - 1;}
        int ext(int d) {x.resize(d + 1);}
        void prt() {
            int l = deg();
            for(register int i = 0; i <= l; ++i)
                printf("%d ", x[i]);
            putchar('\n');
        }
    };
    int RL[maxd], w[maxd], mx2p, N;
    void init(int n) {
        for(N = 1, mx2p = 0; N <= n; N <<= 1, mx2p++);
        for(register int i = 0; i < N; ++i)
            RL[i] = (RL[i >> 1] >> 1) | ((i & 1) << (mx2p - 1));
    }
    void NTT(PLG &A, const int &type) {
        A.ext(N);
        for(register int i = 0; i < N; ++i)
            if(i < RL[i]) swap(A.x[i], A.x[RL[i]]);
        int Wn, x, y;
        for(register int i = 1; i < N; i <<= 1) {
            Wn = fpow(G, (mod - 1) / (i << 1));
            if(type == -1) {Wn = fpow(Wn, mod - 2);}
            w[0] = 1;
            for(register int j = 1; j <= i; ++j)
                w[j] = mul(w[j - 1], Wn);
            for(register int j = 0, p = i << 1; j < N; j += p) {
                for(register int k = 0; k < i; ++k) {
                    x = A.x[j + k], y = mul(A.x[j + k + i], w[k]);
                    A.x[j + k] = (x + y) % mod;
                    A.x[j + k + i] = (x - y + mod) % mod;
                }
            }
        }
    }
    PLG operator *(PLG A, PLG B) {
        int l = A.deg() + B.deg();
        init(l);
        NTT(A, 1), NTT(B, 1);
        for(register int i = 0; i < N; ++i)
            A.x[i] = mul(A.x[i], B.x[i]);
        NTT(A, -1);
        int invN = fpow(N, mod - 2);
        for(register int i = 0; i < N; ++i)
            A.x[i] = mul(A.x[i], invN);
        A.ext(l);
        return A;
    }
    PLG operator -(PLG A, const PLG &B) {
        int l = max(A.deg(), B.deg());
        for(register int i = 0; i <= l; ++i)
            Dec(A.x[i], B.x[i]);
        return A;
    }
}
int step[maxn], inv[maxn];
void pre(int n) {
    step[0] = inv[0] = 1;
    for(register int i = 1; i <= n; ++i) {
        step[i] = mul(step[i - 1], i);
    }
    inv[n] = fpow(step[n], mod - 2);
    for(register int i = n - 1; i >= 1; --i) {
        inv[i] = mul(inv[i + 1], (i + 1));
    }
}
int c[maxn], col[maxn], n, m, u, v;
//sum{c(i)*C(i, k)}
Poly::PLG calc() {
    using namespace Poly;
    PLG A, B, ans;
    A.ext(n), B.ext(n), ans.ext(n);
    for(register int i = 1; i <= n; ++i)
        A.x[i] = mul(c[i], step[i]);
    for(register int i = 1; i <= n; ++i)
        B.x[i] = inv[n - i];
    A = A * B;
    for(register int i = 1; i <= n; ++i)
        ans.x[i] = mul(A.x[i + n], inv[i]);
    return ans;
}
struct edge {
    int v, next;
} e[maxn << 1];
int head[maxn], cnt;
void adde(const int &u, const int &v) {
    e[++cnt] = (edge) {v, head[u]};
    head[u] = cnt;
}
namespace HLD {
    int dfn[maxn], dcnt;
    int son[maxn], top[maxn], fa[maxn], siz[maxn], d[maxn];
    void dfs1(int u, int p) {
        fa[u] = p, siz[u] = 1, d[u] = d[p] + 1;
        for(register int i = head[u]; i; i = e[i].next) {
            int v = e[i].v;
            if(v == p) continue;
            dfs1(v, u), siz[u] += siz[v];
            if(!son[u] || siz[son[u]] < siz[v])
                son[u] = v;
        }
    }
    void dfs2(int u, int tp) {
        top[u] = tp;
        dfn[u] = ++dcnt;
        if(son[u]) dfs2(son[u], tp);
        for(register int i = head[u]; i; i = e[i].next) {
            int v = e[i].v;
            if(v == fa[u] || v == son[u]) continue;
            dfs2(v, v);
        }
    }
    void pre() {dfs1(1, 0), dfs2(1, 1);}
    int LCA(int u, int v) {
        while(top[u] != top[v]) {
            if(d[top[u]] < d[top[v]]) swap(u, v);
            u = fa[top[u]];
        }
        return d[u] < d[v] ? u : v;
    }
    bool cmp_dfn(const int &u, const int &v) {
        return dfn[u] < dfn[v];
    }
    int stk[maxn], stp, cf[maxn], cn[maxn];
    void plus1(int u, int a) {
    	//printf("%d ---- %d\n", u, a);
        while(top[u] != top[a]) {
            ++cf[top[u]];
            if(son[u]) --cf[son[u]];
            u = fa[top[u]];
        }
        if(u == a) return;
        ++cf[son[a]];
        if(son[u]) --cf[son[u]];
    }
    void ins(int x) {
        int lca;
        if(!stp) return stk[++stp] = x, void();
        lca = LCA(x, stk[stp]);
        if(lca == stk[stp]) {
            plus1(x, lca);
            stk[++stp] = x;
            return ;
        }
		while(stp > 1 && dfn[stk[stp]] > dfn[lca])
            --stp;
        if(dfn[lca] < dfn[stk[stp]]) {
            plus1(x, lca), plus1(stk[stp], lca);
            stk[stp] = lca;
        } else {
            plus1(x, lca);
            stk[++stp] = lca;
        }
        stk[++stp] = x;///upd 4.9
    }
    void GetVT(vector<int > &nd) {
        sort(nd.begin(), nd.end(), cmp_dfn);
        int l = nd.size() - 1;
        for(register int i = 0; i <= l; ++i) 
            ins(nd[i]);
        if(!stp) return ;
        ++cn[stk[1]], stp = 0;
    }
    void push(int u, int s) {
        Add(s, cf[u]), cf[u] = s;
        if(son[u]) push(son[u], s);
        for(register int i = head[u]; i; i = e[i].next) {
            int v = e[i].v;
            if(v == fa[u] || v == son[u]) continue;
            push(v, 0);
        }
    }
} 
int main() {
    //freopen("tree.in", "r", stdin);
    using namespace Poly;
    scanf("%d%d", &n, &m);
    for(register int i = 1; i <= n; ++i)
        scanf("%d", &col[i]), nds[col[i]].push_back(i);
    for(register int i = 1; i < n; ++i)
        scanf("%d", &u), adde(i + 1, u), adde(u, i + 1);
    pre(n);
    HLD::pre();
    for(register int i = 1; i <= n; ++i)
        HLD::GetVT(nds[i]);
    HLD::push(1, 0);
    for(register int i = 1; i <= n; ++i)
        ++c[HLD::cf[i] + HLD::cn[i]];
    PLG ans = calc();
    memset(c, 0, sizeof(c));
    for(register int i = 1; i <= n; ++i)
        ++c[HLD::cf[i]];
    ans = ans - calc();
    for(register int i = 1; i <= m; ++i)
        printf("%d ", ans.x[i]);
    return 0;
}
```

Rockdu's Blog

导航

最近发表

大佬们的友链