CF1139C Edgy TreesDFS求连通块大小、思维

题目

题目链接

给一颗树，每条边都是红或黑的。问有多少种大小为$k$的序列$[a_1,a_2...,a_k]$满足，从$a_1$到$a_2$，从$a_2$到$a_3$… …从$a_k-1$到$a_k$，至少经过一条黑边一次。

题解

只需求出所有只包含红边的连通块大小，然后用所有情况减去只在红色连通块中的情况即可。

所有情况即$n^k$。

不包含黑边的情况：求出的第$i$个只包含红边的连通块大小为$S$，那么这个连通块可以有$S^k$个贡献。

因此，答案就是$n^k-S_1^k-S_2^k...-S_m^k$。

代码

/*
 * @Author: hesorchen
 * @Date: 2020-11-26 09:12:46
 * @LastEditTime: 2021-02-04 21:35:52
 * @Description: 栽种绝处的花
 */
#include <bits/stdc++.h>
using namespace std;

const int mod = 1000000007;
long long qpow(long long a, long long b)
{
    long long res = 1;
    while (b)
    {
        if (b & 1)
            res = res * a % mod;
        b /= 2;
        a = a * a % mod;
    }
    return res;
}
struct node
{
    int v, w, next;
} Edge[200010];

int head[100010], ct = 1;
void add(int v, int u, int w)
{
    Edge[ct].v = v;
    Edge[ct].w = w;
    Edge[ct].next = head[u];
    head[u] = ct++;
}
vector<int> siz;
bool vis[100010];

int DFS(int s)
{
    int sum = 1;
    for (int i = head[s]; i; i = Edge[i].next)
    {
        if (!vis[Edge[i].v] && !Edge[i].w)
        {
            vis[Edge[i].v] = 1;
            sum += DFS(Edge[i].v);
        }
    }
    return sum;
}

int main()
{
    int n, k;
    cin >> n >> k;
    for (int i = 1; i < n; i++)
    {
        int u, v, w;
        scanf("%d %d %d", &u, &v, &w);
        add(u, v, w);
        add(v, u, w);
    }
    for (int i = 1; i <= n; i++)
        if (!vis[i])
        {
            vis[i] = 1;
            siz.push_back(DFS(i));
        }
    long long ans = qpow(n, k);
    for (auto it : siz)
        ans = (ans - qpow(it, k) + mod) % mod;
    cout << ans << endl;
    return 0;
}