HDU6035 Colorful Tree

Posted on Edited on In Word count in article: 353 Reading time ≈ 1 mins.
树型Dp

Colorful Tree

简言之, 求总权值减去不合法的权值

即对于每种颜色, 求不经过这种颜色的边的个数

考虑删除这种颜色, 那么剩下的点会构成几个连通块, 对于每个连通块的权值都为\(Size*(Size-1)/2\)

一遍\(DFS\)貌似可以解决了

Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
#include <bits/stdc++.h>

using namespace std;

typedef long long ll;
const int Maxn = 2e5 + 10;
int Head[Maxn], E[Maxn << 1], Next[Maxn << 1];
int N, U, V, Tot, Cur, C[Maxn];
bool Vis[Maxn];
ll Sum[Maxn], Size[Maxn];
ll Res;

inline void AddEdge(int U, int V)
{
    Next[++Cur] = Head[U];
    Head[U] = Cur;
    E[Cur] = V;
}

void DFS(int U, int _F)
{
    Size[U] = 1;
    ll X = Sum[C[U]], Y = X;//这一步很妙啊
    for (int i = Head[U]; i; i = Next[i])
    {
        int V = E[i];
        if (V == _F) continue;
        DFS (V, U);
        Size[U] += Size[V];
        ll Temp = Size[V] - Sum[C[U]] + X;//好好想为啥
        Res += (Temp - 1) * Temp / 2;
        X = Sum[C[U]];
    }
    Sum[C[U]] += Size[U] - (Sum[C[U]] - Y);//简直妙哉
}

inline void Init()
{
    memset (Head, 0, sizeof Head);
    memset (Sum, 0, sizeof Sum);
    memset (Vis, 0, sizeof Vis);
    Tot = Res = Cur = 0;
}

int main()
{
    int Case(0);
    while (scanf ("%d", &N) != EOF)
    {
        Init();
        int Cnt(0);

        for (int i = 1; i <= N; ++i)
        {
            scanf ("%d", C + i);
            if (!Vis[C[i]]) ++Cnt, Vis[C[i]] = 1;
        }

        for (int i = 1; i < N; ++i)
        {
            scanf ("%d %d", &U, &V);
            AddEdge(U, V), AddEdge(V, U);
        }

        DFS(1, 0);

        ll Ans = 1ll * N * (N - 1) * Cnt / 2;

        for (int i = 1; i <= N; ++i)
            if (Vis[i])
            {
                ll Temp = N - Sum[i];
                Ans -= Temp * (Temp - 1) / 2;//这个是最后一个连通块
            } 
        
        printf ("Case #%d: %lld\n", ++Case, Ans - Res);
    }
    system("pause");
}