CodeForces - 1139C Edgy Trees (快速幂+dfs）

题目：CodeForces - 1139C 

题意：一个n个节点的无向图，有n-1条边每条边是黑色或者红色，现在统计包含k个点（不需要连续）的路中至少有1条黑边的路径数目

分析：总共有n^k条路径，将所有通过红边连在一起的联通块的数目，同时记录每个联通块中点的数目p[i]；

由于这个可以自反（就是4可以走（4->4））所以每个联通块中的没有黑色边的路径数目为p[i] ^ k;

所以总共的数目就是 % mod;

但是需要注意的是如果减出来的差小于0，应该加上mod;没考虑到这个就一直wa

#include <iostream>
#include <stdio.h>
#include <string.h>
#include <vector>
#include <algorithm>

using namespace std;

typedef long long LL;
const int maxn = 1e5 + 5;
const int mod = 1e9 + 7;
bool vis[maxn];

struct Edge
{
    int to;
    int color;
    Edge(int _to,int _color):to(_to),color(_color){}
};

LL p[maxn];//record the number of the connected block
vector<Edge>edge[maxn];
void dfs(int x,int cnt)
{
    vis[x] = 1;
    int sz = edge[x].size();
    for(int i = 0; i < sz;i ++)
    {
        Edge v = edge[x][i];
        if(v.color == 0 && !vis[v.to])
        {
            p[cnt] += 1;
            dfs(v.to,cnt);
        }
    }
}
LL pow_mod(LL a, LL n, LL p){
    LL ret = 1;
    LL tmp = a;
    while(n) {
        if(n & 1) ret = (ret * tmp) % p;
        tmp = tmp * tmp % p;
        n >>= 1;
    }
    return ret;
}

int main()
{
    LL n,k;
    while(cin >> n >> k)
    {
        int u,v,x;
        for(int i = 1;i < n;i ++)
        {
            cin >> u >> v >> x;
            edge[u].push_back(Edge(v,x));
            edge[v].push_back(Edge(u,x));
        }
        memset(vis,false,sizeof(vis));
        int cnt = 1;
        for(int i = 1;i <= n;i ++)
        {
            p[i] = 1;
        }
        for(int i = 1;i <= n;i ++)
        {
            if(!vis[i])
            {
                dfs(i,cnt++);
            }
        }
        LL sum = pow_mod(n,k,mod);
        LL ans = 0;
        for(int i = 1;i < cnt;i ++)
        {
            ans = (ans + pow_mod(p[i],k,mod)) % mod;
        }
        ans = sum - ans;
        if(ans < 0) ans += mod;
        cout << ans << endl;
    }
    return 0;
}