HDU6035 Colorful Tree(树形dp)

最新推荐文章于 2021-08-13 16:47:02 发布

原创最新推荐文章于 2021-08-13 16:47:02 发布 · 1.7k 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#算法 #数据结构 #动态规划

算法同时被 3 个专栏收录

246 篇文章

订阅专栏

动态规划

62 篇文章

订阅专栏

数据结构

46 篇文章

订阅专栏

本文介绍了一种算法，用于计算一棵带有颜色标记节点的树的所有路径中不同颜色的数量总和。通过一次深度优先搜索（DFS），巧妙地计算每种颜色对路径总数的贡献。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

Colorful Tree

Time Limit: 6000/3000 MS (Java/Others) Memory Limit: 131072/131072 K (Java/Others)
Total Submission(s): 1152 Accepted Submission(s): 463

Problem Description

There is a tree with

n nodes, each of which has a type of color represented by an integer, where the color of node

i is

ci .

The path between each two different nodes is unique, of which we define the value as the number of different colors appearing in it.

Calculate the sum of values of all paths on the tree that has

n(n−1)2 paths in total.

Input

The input contains multiple test cases.

For each test case, the first line contains one positive integers

n , indicating the number of node.

(2≤n≤200000)

Next line contains

n integers where the

i -th integer represents

ci , the color of node

i .

(1≤ci≤n)

Each of the next

n−1 lines contains two positive integers

x,y

(1≤x,y≤n,x≠y) , meaning an edge between node

x and node

y .

It is guaranteed that these edges form a tree.

Output

For each test case, output " Case #

x :

y " in one line (without quotes), where

x indicates the case number starting from

1 and

y denotes the answer of corresponding case.

Sample Input

      
       3
1 2 1
1 2
2 3
6
1 2 1 3 2 1
1 2
1 3
2 4
2 5
3 6

Sample Output

      
       Case #1: 6
Case #2: 29

Source

2017 Multi-University Training Contest - Team 1

Recommend

liuyiding | We have carefully selected several similar problems for you: 6044 6043 6042 6041 6040

题意:给一棵数,每个点都有一种不同的颜色表示.题目要求求出对于每一条路径,不经过的颜色数的总和.

首先把问题转化成,求每种颜色的贡献,每种颜色的贡献就是经过该种颜色的不同路径数目的和,反向思考一下,就是路径总数减去不经过该种颜色的路径数.

我们考虑某个节点u,它的颜色为c,儿子个数为n,如果我们现在知道了,在节点u的儿子中,m个节点的颜色也为c,他们的儿子个数的总和为s,那么该节点的贡献是不是就是C(n-s,2)

我们需要在一遍dfs中解决这个问题

记录两个数组,一个是son,代表每个节点的孩子总数,一个是sum,记录对于每种颜色,当前已经处理过的该颜色的节点的孩子总数.

#include <cstring>
#include <iostream>
#include <stdio.h>
#include <algorithm>
#include <cmath>
#include <map>
#include<time.h>
using namespace std;
typedef long long ll;
const int MAXN=200000+10;
struct edge{
    int v,next;
}e[MAXN<<1];
int tot=0;
int head[MAXN];
int c[MAXN];
ll sum[MAXN];
ll son[MAXN];
ll res=0;
void add(int u,int v){
    e[tot].v=v;
    e[tot].next=head[u];
    head[u]=tot++;
}
void dfs(int u,int p){
    son[u]=1;
    ll y=sum[c[u]];
    ll x=sum[c[u]];
    for(int k=head[u];k!=-1;k=e[k].next){
        int v=e[k].v;
        if(v!=p){
            dfs(v,u);
            son[u]+=son[v];
            ll temp=son[v]-sum[c[u]]+x;
            res+=(temp-1)*temp/2;
            x=sum[c[u]];
        }
    }
    sum[c[u]]+=son[u]-(sum[c[u]]-y);
}

int vis[MAXN];
void init(){
    memset(head,-1,sizeof(head));
    memset(sum,0,sizeof(sum));
    memset(vis,0,sizeof(vis));
    tot=0;
    res=0;
}
int main(){
    int cas=1;
    int n;
    //freopen("1003.in","r",stdin);
    while(scanf("%d",&n)!=EOF){
        ll cnt=0;
        init();
        for(int i=1;i<=n;i++){
            scanf("%d",c+i);
            if(!vis[c[i]]){
                cnt++;
                vis[c[i]]=1;
            }
        }
        for(int i=0;i<n-1;i++){
            int u,v;
            scanf("%d%d",&u,&v);
            add(u,v);
            add(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",cas++,ans-res);
    }
}