J - Jingle Bells Gym - 102448J（樹形dp）

Christmas is coming, and Kinho was building his christmas tree! As a good programmer, his tree is, of course, an acyclic, undirected and connected graph, with N vertices and N−1 unweighted edges.

Kinho also has 5 buckets of ink, with colors {1, 2, 3, 4, 5}, having an infinite amount of ink on each bucket. The absence of color is indicated with the number 0. Initially, all edges of the tree had no color, but then he decided to paint some edges with one of the five available colors.

Kinho quickly got tired and asked for Tfg’s help. He would need to continue painting the tree, in a way that it became beautiful.

A tree is beautiful if all of its edges are painted, and there is no vertex with two or more adjacent edges with the same color.

Tfg likes beautiful trees, and then, in order to raise is excitement, he asked you how many ways are there to paint the entire tree. Two ways are considered different if they have at least one edge with a different color in each way.

Note that there will be edges which have been already painted by Kinho - and Tfg must not change them! He is only allowed to paint uncolored edges.

Input
The first line of the input contains a single integer N, 1≤N≤105, the number of vertices in the tree. Then, N−1 lines follow, each of them containing three integers u, v and c, (1≤u,v≤N), (0≤c≤5), describing that there exists an edge between vertices u and v with color c. Remember that if c=0, then the edge is uncolored.

Output
You must output a single integer, the number of beautiful paintings, given the initial description of the tree. Since the number can be very large, output its remainder when divided by 109+7.

Examples
Input
4
1 2 0
1 3 1
3 4 0
Output
16
Input
4
1 2 0
2 3 0
3 4 0
Output
80

題意：
有5種顏色。相鄰邊不能染同一種顏色，有些邊已經染顏色，求染色方案數。

思路：
定義$f[i][j]$爲以$i$爲根節點，且i的邊的選擇方案爲$j$的方案數。

#include <cstdio>
#include <cstring>
#include <algorithm>

using namespace std;

typedef long long ll;
const int maxn = 2e5 + 7;
const int mod = 1e9 + 7;

ll f[maxn][35];
int head[maxn],nex[maxn],to[maxn],val[maxn],tot;

void add(int x,int y,int z) {
    to[++tot] = y;
    nex[tot] = head[x];
    val[tot] = z;
    head[x] = tot;
}

void dfs(int u,int fa) {
    f[u][0] = 1;
    for(int i = head[u];i;i = nex[i]) {
        int v = to[i],w = val[i];
        if(v == fa) continue;
        int l = 0,r = 4;
        if(w) l = r = w - 1;
        dfs(v,u);
        for(int j = (1 << 5) - 1;j >= 0;j--) {
            for(int num = l;num <= r;num++) {
                if(f[u][j] == 0 || (j & (1 << num))) continue;
                for(int k = 0;k < (1 << 5);k++) {
                    if(k & (1 << num)) continue;
                    f[u][j | (1 << num)] += f[u][j] * f[v][k] % mod;
                    f[u][j | (1 << num)] %= mod;
                }
            }
            f[u][j] = 0;
        }
    }
}

int main() {
    int n;scanf("%d",&n);
    for(int i = 1;i < n;i++) {
        int x,y,z;scanf("%d%d%d",&x,&y,&z);
        add(x,y,z);add(y,x,z);
    }
    dfs(1,-1);
    ll ans = 0;
    for(int i = 0;i < (1 << 5);i++) {
        ans += f[1][i];
        ans %= mod;
    }
    printf("%lld\n",ans);
    return 0;
}