HDU6035 Colorful Tree

树形DP与线段树合并求解颜色路径问题

最新推荐文章于 2019-02-15 18:01:00 发布

原创最新推荐文章于 2019-02-15 18:01:00 发布 · 1.7k 阅读

1 ·

CC 4.0 BY-SA版权

多校同时被 2 个专栏收录

9 篇文章

订阅专栏

线段树合并

3 篇文章

订阅专栏

本文介绍了一种使用树形DP和线段树合并技术来解决特定图论问题的方法，即计算一棵树中所有路径上不同颜色节点数量之和。通过递归地遍历树结构，并利用线段树来高效更新和查询颜色分布，最终得到所有路径上的颜色数总和。

计算所有路径上不同颜色数的和

假设每种颜色对每条路径都有贡献，再减去对于每种颜色路径没有贡献的情况。（一开始觉得答案很大自信取模。。。）

O（nlogn）:树形DP+线段树合并

#include<stdio.h>
#include<math.h>
#include<string.h>
#include<stdlib.h>
#include<string>
#include<iostream>
#include<algorithm>
using namespace std;
typedef long long ll;
const int MAXN=(int)2e5+10;
const int MOD=(int)1e9+7;
struct node{
    int to,nxt;
}ed[MAXN<<1];
int head[MAXN],cnt;
int root[MAXN];
int ls[MAXN*20],rs[MAXN*20],sum[MAXN*20],tol;
void insert(int &rt,int l,int r,int x,int y){
    if(rt==0)rt=++tol;
    if(l==r){
        sum[rt]=y;
        return ;
    }
    int mid=l+r>>1;
    if(x<=mid)insert(ls[rt],l,mid,x,y);
    else insert(rs[rt],mid+1,r,x,y);
    sum[rt]=sum[ls[rt]]+sum[rs[rt]];
}
int query(int rt,int l,int r,int k){
    if(rt==0)return 0;
    if(l==r)return sum[rt];
    int mid=l+r>>1;
    if(k<=mid)return query(ls[rt],l,mid,k);
    else return query(rs[rt],mid+1,r,k);
}
int merge(int u,int v){
    if(u==0||v==0)return u|v;
    ls[u]=merge(ls[u],ls[v]);
    rs[u]=merge(rs[u],rs[v]);
    sum[u]+=sum[v];
    return u;
}
void addedge(int u,int v){
    ed[cnt].to=v;
    ed[cnt].nxt=head[u];
    head[u]=cnt++;
}
int col[MAXN],sz[MAXN],vis[MAXN];
int n;
ll de,ans;
void dfs(int u,int pre){
    sz[u]=1;
    int step=0;
    for(int i=head[u];i!=-1;i=ed[i].nxt){
        int v=ed[i].to;
        if(v!=pre){
            dfs(v,u);
            sz[u]+=sz[v];
            ll qr=query(root[v],1,n,col[u]);
            ll szv=sz[v]-qr;
            de=(de+1LL*szv*(szv-1)/2);
            step+=qr;    
        }
    }    
    insert(root[u],1,n,col[u],sz[u]-step);
    for(int i=head[u];i!=-1;i=ed[i].nxt){
        int v=ed[i].to;
        if(v!=pre){
            merge(root[u],root[v]);
        }
    }
}
int main()
{
    int ca=1;
    while(~scanf("%d",&n)){
        de=0;
        tol=0;
        cnt=0;
        for(int i=1;i<=n;i++)root[i]=0;
        memset(ls,0,sizeof(ls));
        memset(rs,0,sizeof(rs));
        memset(vis,0,sizeof(vis));
        for(int i=1;i<=n;i++)head[i]=-1;
        for(int i=1;i<=n;i++)scanf("%d",&col[i]),vis[col[i]]=1;
        for(int i=1;i<n;i++){
            int u,v;
            scanf("%d%d",&u,&v);
            addedge(u,v);
            addedge(v,u);
        }
        dfs(1,0);
        for(int i=1;i<=n;i++){
            if(vis[i]&&i!=col[1]){
                ll qr=query(root[1],1,n,i);
                ll szv=n-qr;
                de=(de+1LL*szv*(szv-1)/2);
            }
        }
        int colnum=0;
        for(int i=1;i<=n;i++)colnum+=vis[i];
        ans=1LL*n*(n-1)/2*colnum;
        ans=((ans-de));
        printf("Case #%d: %lld\n",ca++,ans);
    }
    return 0;
}

O（n）:

#include<stdio.h>
#include<math.h>
#include<string.h>
#include<stdlib.h>
#include<string>
#include<iostream>
#include<algorithm>
using namespace std;
typedef long long ll;
const int MAXN=(int)2e5+10;
const int MOD=(int)1e9+7;
struct node{
    int to,nxt;
}ed[MAXN<<1];
int head[MAXN],cnt;
int root[MAXN];
int sum[MAXN];
void addedge(int u,int v){
    ed[cnt].to=v;
    ed[cnt].nxt=head[u];
    head[u]=cnt++;
}
int col[MAXN],sz[MAXN],vis[MAXN];
int n;
ll de,ans;
void dfs(int u,int pre){
    sz[u]=1;
    int step=0;
    int s1=sum[col[u]];
    for(int i=head[u];i!=-1;i=ed[i].nxt){
        int v=ed[i].to;
        if(v!=pre){
            dfs(v,u);
            sz[u]+=sz[v];
            int qr=sum[col[u]]-s1;
            int szv=sz[v]-qr;
            s1=sum[col[u]];
            de=(de+1LL*szv*(szv-1)/2);    
        	step+=qr;
		}
    }    
    sum[col[u]]+=sz[u]-step;
}
int main()
{
    int ca=1;
    while(~scanf("%d",&n)){
        de=0;
        cnt=0;
        for(int i=1;i<=n;i++)vis[i]=0;
        for(int i=1;i<=n;i++)head[i]=-1;
        for(int i=1;i<=n;i++)scanf("%d",&col[i]),vis[col[i]]=1,sum[col[i]]=0;
        for(int i=1;i<n;i++){
            int u,v;
            scanf("%d%d",&u,&v);
            addedge(u,v);
            addedge(v,u);
        }
        dfs(1,0);
        for(int i=1;i<=n;i++){
            if(vis[i]&&i!=col[1]){
                int qr=sum[i];
                int szv=n-qr;
                de=(de+1LL*szv*(szv-1)/2);
            }
        }
        int colnum=0;
        for(int i=1;i<=n;i++)colnum+=vis[i];
        ans=1LL*n*(n-1)/2*colnum;
        ans=((ans-de));
        printf("Case #%d: %lld\n",ca++,ans);
    }
    return 0;
}