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

71 阅读2分钟

题目

题目链接

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

题解

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

所有情况即 n k n^k nk。

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

因此,答案就是 n k − S 1 k − S 2 k . . . − S m k n^k-S_1^k-S_2^k...-S_m^k nk−S1k​−S2k​...−Smk​。

代码

/*
 * @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;
}