牛客竞赛题目讲解_Tree Intersection(树上启发式合并)

//  https://ac.nowcoder.com/acm/contest/1112/I

#include "stdafx.h"

//#include <bits/stdc++.h>

#include <iostream>

#include <cstdio>

#include <cstring>

#include <algorithm>

 

#include <vector> 

 

using namespace std; 

const int maxn = 1e5 + 5;

 int n, c[maxn]={0,

1,1,2,2 ,2,3,3,5

};

int edge[32][2]={

1,7,

7,8,

1 ,2,

2, 3,

3 ,4,

4,5,

5,6

};

  

int sum[maxn], fa[maxn];

int sz[maxn], son[maxn];

int dfsn[maxn],T=0;

int a[maxn];//dfs序的颜色

vector<int>load[maxn];

 

void dfs1(int s, int pre) {

    sz[s] = 1;// 子树大小

    fa[s] = pre;

dfsn[s]=++T; //dfs序 编号

    for(auto e : load[s]) {

        if(e == pre)

            continue;

        dfs1(e, s);

        sz[s] += sz[e];

        if(sz[e] > sz[son[s]])

            son[s] = e;// 重子树

    }

}

int ans[maxn], cnt[maxn],  last, first;

  

void dfs2(int s, int pre, int flag) {

    for(auto e : load[s]) {

        if(e != pre && e != son[s])// 先 计算轻链

        dfs2(e, s, 0);

    }

    if(son[s])

        dfs2(son[s], s, 1);// 最后计算重链，而不需要清除 cnt[maxn],  last, first;

    

for(auto e : load[s]) 

   if(e != pre && e != son[s])// 只 计算轻链

{  

  for(int i = 0; i <sz[e]; i++) {

     cnt[a[dfsn[e]+i]] ++;

   if(cnt[a[dfsn[e]+i]] == 1)// 有这种颜色节点在子树上

{

            last++;

cout<<"last: e="<<e<<", val="<<"1"<<endl;

}

   

        if(cnt[a[dfsn[e]+i]] == sum[a[dfsn[e]+i]]) // 全部(有第一次出现的)这种颜色节点在子树上 

{

            first++;

cout<<"first:   e="<<e<<", val="<<"1"<<endl;

}

   

  }

}

    cnt[c[s]] ++;

 

   if(cnt[c[s]]  == 1)// 有这种颜色节点在子树上

{

            last++;

cout<<"last: s="<<s<<", val="<<"1"<<endl;

}

        if(cnt[c[s]]  == sum[c[s]]) // 全部(有第一次出现的)这种颜色节点在子树上 

{

            first++;

cout<<"first:   s="<<s<<", val="<<"1"<<endl;

}

    ans[s] = last - first;

cout<<"  s="<<s<<", first="<<first<<", last="<<last<<endl;

   if(!flag) 

{// 轻链 clear

         for(int i = 0; i <sz[s]; i++) {

     cnt[a[dfsn[s]+i]] --;

   if(cnt[a[dfsn[s]+i]] == 0)// 有这种颜色节点在子树上

{

            last--;

cout<<"last: s="<<s<<", val="<<"-1"<<endl;

}

        if(cnt[a[dfsn[s]+i]] == sum[a[dfsn[s]+i]]-1) // 全部(有第一次出现的)这种颜色节点在子树上 

{

            first--;

cout<<"first:   s="<<s<<", val="<<"-1"<<endl;

}    

  } 

    }

}

int s[maxn], e[maxn];

void init() {

    for(int i = 1; i <= n; i++) {

        load[i].clear();

        cnt[i] = sum[i] = son[i] = 0;

T=0;

    }

}

int main() 

{

n=8;

    //while(cin >> n) 

{

        init();

        for(int i = 1; i <= n; i++) {

          //  cin >> c[i];

            sum[c[i]]++;//颜色统计

        }

        for(int i = 1; i < n; i++) {

            //cin >> s[i] >> e[i];

s[i]=edge[i-1][0];

e[i]=edge[i-1][1];

            load[s[i]].push_back(e[i]);

            load[e[i]].push_back(s[i]);

        }

        dfs1(1, 0);

  for(int i = 1; i <=n; i++) cout<<"i=  "<<i<<":   son[i]="<<son[i]<<endl;

    for(int i=1; i<=n; i++){ 

a[dfsn[i]]=c[i];

}

        dfs2(1, 0, 1);

cout<<"The end: first="<<first<<", last="<<last<<endl;

        for(int i = 1; i < n; i++) {

            if(fa[s[i]] == e[i])

                cout << ans[s[i]] << endl;

            else

                cout << ans[e[i]] << endl;

        }

    }

}