编程之美2015初赛第一场: 彩色的树

题目1 : 彩色的树

时间限制: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

解答思路

本题主要解决任意时刻树中有多少棵颜色相同的子树，可以记录每一个节点的邻居节点，在每次更新一个节点的颜色的时候，都重新计算有多少棵同色子树。neighbor[i]表示节点i的邻居节点集合，color[i]表示节点i的颜色，new_color表示节点i新获得的颜色，color_tree_num表示树中有多少棵同色子树，初始为1。如果color[i] == new_color，color_tree_num不变；如果color[i] != new_color,节点获得了新颜色，需要更新color_tree_num　　for(节点i的邻居节点 p: neighbor[i]){　　　　if(p.color == color[i] /*&& p.color != new_color*/) color_tree_num++;　　　　else if(/*p.color != color[i] &&*/ p.color == new_color) color_tree_num--;　　　　/*else if(p.color != color[i] && p.color != new_color) color_tree_num不变*/　　}

代码如下

#include<iostream>
#include<cstring>
#include<cstdio>
#include<unordered_set>
using namespace std;
int color[100010];
struct Node{
    //int i;
    int n;
    int *p;
};
Node neighbor[100010];
int main(){
    int T, n, q, i, j, query, c;
    unordered_set<long long> edge;
    memset(color, 0, sizeof(color));
    scanf("%d", &T);
    for(int k = 1; k <= T; k++){
        int tree = 1;
        printf("Case #%d:\n", k);
        scanf("%d", &n);
        for(int t = 1; t < n; t++){
            scanf("%d%d", &i, &j);
            color[i]++;
            color[j]++;
            long long ii = i, jj = j;
            if(i < j)
                edge.insert((ii << 32) | jj);
            else
                edge.insert((jj << 32) | ii);
        }
        for(i = 1; i <= n; i++){
            neighbor[i].n = 0;
            neighbor[i].p = new int[color[i]];
        }
        for(auto it = edge.begin(); it != edge.end(); ++it){
            j = int(*it);
            i = int((*it) >> 32);
            neighbor[i].p[neighbor[i].n++] = j;
            neighbor[j].p[neighbor[j].n++] = i;
        }

        memset(color, 0, sizeof(color));
        scanf("%d", &q);
        for(int t = 0; t < q; t++){
            scanf("%d", &query);
            if(query == 1){
                printf("%d\n", tree);
            }else{//query == 2
                scanf("%d%d", &i, &c);
                if(color[i] != c){
                    for(int* nei = neighbor[i].p, nn = neighbor[i].n, m = 0; m < nn; m++){
                        j = nei[m];
                        if(color[i] == color[j]){
                            tree++;
                        }else if(color[j] == c){
                            tree--;
                        }
                    }
                    color[i] = c;
                }
            }
        }
    }
    return 0;
}