HDU 6458 Checkout

题意：

Checkout

Time Limit: 2000/1000 MS (Java/Others)    Memory Limit: 524288/524288 K (Java/Others)
Total Submission(s): 35    Accepted Submission(s): 16

 

Problem Description

Alice 是一个身怀改变世界的抱负的著名企业家，手中掌控很多著名的公司，为了更好的管理， Alice 建立了一套很完善的架构体系，已知 Alice 的企业的架构体系是一棵树，每个节点代表一个人。对于每个节点，它的父节点就是这个人的直接 leader，每个节点都有一个权值，代表这个人的爱好，每对属于同一直接 leader 的节点如果有相同的爱好那么他们就会给公司产生1的和谐度，如果一个人和他/她的直接 leader 的爱好相同也会给公司产生1的和谐度，但是再完美的公司都会产生离职的情况，如果一个人离职，并且这个人有直接下属，就从这个人的直接下属中选一个成为这个部门新的 leader，使得新的组织架构的和谐度最高，新晋的人汇报给离职的人的 leader（如果存在的话），同时新 leader 的原来下属的直接 leader 不变， Alice 想知道如果某个人离职整个公司的和谐度是多少。

 

 

Input

第一行两个正整数 n, m 分别代表公司的总人数和爱好的种数。
第二行有 n 个整数，第 i 个数 ai，代表第 i 个人的爱好。
后面 n − 1 行每行两个数u，v代表 u，v 存在上下属关系（上下级关系不确定）。
假设树根为 1，即 1 号员工是公司的ceo，他/她不需要汇报给任何人。
1 ≤ m ≤ n ≤ 100, 000
1 ≤ ai ≤ m

 

 

Output

输出一行n个数，第i个数代表如果编号为 i 的人离职，公司的和谐度会是多少，以空格分隔，行末无空格。

 

 

Sample Input

4 3 2 3 1 3 1 2 2 3 2 4

 

 

Sample Output

1 0 1 0

 

 

Source

755hdu女生赛

分析：先计算出一棵树总的和谐度，再枚举每个子节点代替父节点时和谐度的变化，取最大值，两遍dfs就好了，计数细节都在代码里面。

代码：

#include<bits/stdc++.h>
#define ll long long
using namespace std;
const int N = 1e5+5;
map<int,int> mp[N];
int n,m,u,v,a[N],f[N];
vector<int> g[N];
ll sum, ans[N];
void dfs(int u, int fa) {
    f[u] = fa;
    for(int i=0;i<g[u].size();i++) {
        int v=g[u][i];
        if(v!=fa) {
            dfs(v, u);
            if (a[v] == a[u]) sum++;
            sum += mp[u][a[v]];
            mp[u][a[v]]++;
        }
    }
}

void DFS(int u) {
    int rt = f[u];
    if (rt) {
        ll res = sum-mp[rt][a[rt]];
        mp[rt][a[u]]--;///去掉u作为兄弟节点的贡献
        for (auto tp : mp[u])///去掉兄弟节点贡献
            res += 1ll * tp.second * mp[rt][tp.first];
        mp[rt][a[u]]++;///恢复
        ans[rt] = max(ans[rt], res);
    }
    if (g[u].size()<2 && u!=1) {///根节点特判
        ans[u] = sum - mp[rt][a[u]] + 1;///去掉父子节点的贡献
        if (a[rt] == a[u])///去掉兄弟节点的贡献
            ans[u]--;
    }
    for(int i=0;i<g[u].size();i++) {
        int v=g[u][i];
        if(v!=f[u]) DFS(v);
    }
}
int main() {
    scanf("%d%d",&n,&m);
    for (int i = 1; i <= n; i++) scanf("%d",&a[i]);
    for (int i = 1; i < n; i++) {
        scanf("%d%d",&u,&v);
        g[u].push_back(v);
        g[v].push_back(u);
    }
    dfs(1, 0);
    DFS(1);
    for (int i = 1; i < n; i++) printf("%lld ", ans[i]);
    printf("%lld\n",ans[n]);
    return 0;
}