多彩的树 ssl2658 dfs+容斥原理

Description

---

我们将给出一颗有N个节点的树，每个节点都有一个颜色，你可以任意选择一个节点作为起点，并通向它所能到达的节点，然后最终停在一个终点上(起点和终点可能相同)，当然了每一个结点只能经过一次,叶子的颜色种类为K。请问有多少种不同的方案可以使得你经过的节点的颜色正好为K种？(两种方案不同当且仅当起点不同或终点不同)。

Solution

---

我果然还是太弱了

k只有10，那么考虑用容斥原理。考虑去掉颜色的限制，那么一棵n节点的树总路径为n2 的，那么路径方案数即枚举哪些颜色不被选，剩下其它颜色可以选或不选。然后把这些颜色对应的点全部删去，得到一个森林，接下来O(n)统计答案，最后乘上容斥系数即可。
复杂度O(2k∗n)

然后还可以点分治的，但是不会

我果然还是太弱了

Code

---

#include <stdio.h>#include <string.h>#include <vector>#define rep(i, st, ed) for (int i = st; i <= ed; i += 1)#define erg(i, st) for (int i = ls[st]; i; i = e[i].next)#define fill(x, t) memset(x, t, sizeof(x))#define ll long long#define N 50001#define E N << 1 | 1struct edge{int x, y, next;}e[E];int ls[N];inline void addEdge(int &cnt, int x, int y){    cnt += 1; e[cnt] = (edge){x, y, ls[x]}; ls[x] = cnt;    cnt += 1; e[cnt] = (edge){y, x, ls[y]}; ls[y] = cnt;}int vis[N], c[N];int size;inline void getAns(int now, int col){    vis[now] = 1;    size += 1;    erg(i, now){        if (!vis[e[i].y] && !(col & c[e[i].y])){            getAns(e[i].y, col);        }    }}using std:: vector;vector<int> sch, poe;inline void dfsCol(int foe, int dep, int col, int m){    if (dep == m){        sch.push_back(col);        poe.push_back(foe);    }else{        dfsCol(foe, dep + 1, col, m);        dfsCol(-foe, dep + 1, col | (1 << dep), m);    }}int main(void){    int n, m;    scanf("%d%d", &n, &m);    rep(i, 1, n){        int tmp;        scanf("%d", &tmp);        c[i] = 1 << (tmp - 1);    }    int edgeCnt = 0;    rep(i, 2, n){        int x, y;        scanf("%d%d", &x, &y);        addEdge(edgeCnt, x, y);    }    dfsCol(1, 0, 0, m);    ll ans = 0;    rep(i, 0, sch.size() - 1){        fill(vis, 0);        rep(st, 1, n){            if (!vis[st] && !(sch[i] & c[st])){                size = 0;                getAns(st, sch[i]);                ans += (ll) (size * size * poe[i]);                // printf("%d %d\n", size * size * poe[i], sch[i]);            }        }    }    printf("%lld\n", ans);    return 0;}