多彩的树 ssl2658 dfs+容斥原理

Description

---

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

Solution

---

我果然还是太弱了

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