[2015编程之美] 第一场A

#1156 : 彩色的树

时间限制: 2000ms

单点时限: 1000ms

内存限制: 256MB

描述

给定一棵n个节点的树，节点编号为1, 2, …, n。树中有n - 1条边，任意两个节点间恰好有一条路径。这是一棵彩色的树，每个节点恰好可以染一种颜色。初始时，所有节点的颜色都为0。现在需要实现两种操作:

1. 改变节点x的颜色为y；

2. 询问整棵树被划分成了多少棵颜色相同的子树。即每棵子树内的节点颜色都相同，而相邻子树的颜色不同。

输入

第一行一个整数T，表示数据组数，以下是T组数据。

每组数据第一行是n，表示树的节点个数。接下来n - 1行每行两个数i和j，表示节点i和j间有一条边。接下来是一个数q，表示操作数。之后q行，每行表示以下两种操作之一：

1. 若为"1"，则询问划分的子树个数。

2. 若为"2 x y"，则将节点x的颜色改为y。

输出

每组数据的第一行为"Case #X:"，X为测试数据编号，从1开始。

接下来的每一行，对于每一个询问，输出一个整数，为划分成的子树个数。

数据范围

1 ≤ T ≤ 20

0 ≤ y ≤ 100000

小数据

1 ≤ n, q ≤ 5000

大数据

1 ≤ n, q ≤ 100000

样例输入

2

3

1 2

2 3

3

1

2 2 1

1

5

1 2

2 3

2 4

2 5

4

1

2 2 1

2 3 2

1

样例输出

Case #1:

1

3

Case #2:

1

5

TLE代码，能过小数据。

#include <iostream>

#include <cstdio>

#include <cstring>

using namespace std;

#define ll long long

#define N 100010



struct Edge

{

    int to,next;

}edge[N];

int tot;

int col[N];

int head[N];

int n,m;

int ans;



void init()

{

    tot=0;

    ans=1;

    for(int i=1;i<=n;i++) col[i]=0;

    memset(head,-1,sizeof(head));

}

void add(int u,int v)

{

    edge[tot].to=v;

    edge[tot].next=head[u];

    head[u]=tot++;

}

void update(int u,int c)

{

    for(int i=head[u];i!=-1;i=edge[i].next)

    {

        int v=edge[i].to;

        if(col[v]!=col[u] && col[v]==c) ans--;

        else if(col[v]==col[u] && col[v]!=c) ans++;

    }

    col[u]=c;

}

int main()

{

    int T,iCase=1;

    scanf("%d",&T);

    while(T--)

    {

        init();

        scanf("%d",&n);

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

        {

            int u,v;

            scanf("%d%d",&u,&v);

            add(u,v);

            add(v,u);

        }

        printf("Case #%d:\n",iCase++);

        scanf("%d",&m);

        while(m--)

        {

            int op,pos,c;

            scanf("%d",&op);

            if(op==1) printf("%d\n",ans);

            else

            {

                scanf("%d%d",&pos,&c);

                update(pos,c);

            }

        }

    }

    return 0;

}

 

正解代码：

#include <iostream>

#include <cstdio>

#include <map>

#include <cstring>

using namespace std;

#define ll long long

#define N 100010



struct Edge

{

    int to,next;

}edge[N<<1];



int n,m;

int ans;

int tot;

int fa[N];

int vis[N];

int col[N];

int head[N];

map<int,int> mp[N];



void init()

{

    tot=0;

    ans=1;

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

    {

        col[i]=0;

        mp[i].clear();

    }

    memset(vis,0,sizeof(vis));

    memset(head,-1,sizeof(head));

}

void add(int u,int v)

{

    edge[tot].to=v;

    edge[tot].next=head[u];

    head[u]=tot++;

}

void dfs(int u)

{

    vis[u]=1;

    for(int i=head[u];i!=-1;i=edge[i].next)

    {

        int v=edge[i].to;

        if(!vis[v])

        {

            fa[v]=u;

            mp[u][0]++;

            dfs(v);

        }

    }

}

void update(int x,int c)

{

    if(col[x]==c) return;

    int y=fa[x];

    //对儿子节点

    ans+=mp[x][col[x]];

    ans-=mp[x][c];

    //对父亲节点

    if(y!=-1)

    {

        if(c!=col[y]) ans++;

        if(col[x]!=col[y]) ans--;

        mp[y][col[x]]--;

        mp[y][c]++;

    }

    col[x]=c;

}

int main()

{

    int T,iCase=1;

    scanf("%d",&T);

    while(T--)

    {

        init();

        scanf("%d",&n);

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

        {

            int u,v;

            scanf("%d%d",&u,&v);

            add(u,v);

            add(v,u);

        }

        fa[1]=-1;

        dfs(1);

        printf("Case #%d:\n",iCase++);

        scanf("%d",&m);

        while(m--)

        {

            int op,pos,c;

            scanf("%d",&op);

            if(op==1) printf("d\n",ans);

            else

            {

                scanf("%d%d",&pos,&c);

                update(pos,c);

            }

        }

    }

    return 0;

}