[牛客竞赛]生活在树上（图论 + 树）

题目传送门：生活在树上 (nowcoder.com)

题目描述

请看原题。

这里介绍两种思路。

思路一

将点 i 当作根，父节点、子节点都当作子节点。

只要建立一个双向边的图就可以完成。相当于先把树建立起来，再把某个点i拉出来作为根，那么这又可以形成一颗新的以i为根的树。

那么每个点能够到达的点就是每个点的孙子节点。

在输入的时间将d = 2的边直接加到答案中，d = 1的记录边，d > 2的全部忽略。

代码一（449ms）

#include <bits/stdc++.h>
#define pb push_back
#define int long long
using namespace std;

inline int readInt()
{
    int s = 0, w = 1;char ch = getchar();
    for(;ch > '9' || ch < '0';ch = getchar())if(ch == '-')w = -1;
    for(;'0' <= ch && ch <= '9';ch = getchar())s = s * 10 + ch - '0';
    return s * w;
}

const int maxn = 1e6 + 10;
int ans[maxn];
vector<int> g[maxn];


signed main()
{
    ios::sync_with_stdio(0);
	int N = readInt();
    for(int i = 2;i <= N; ++ i)
    {
        int x = readInt(), w = readInt();
        if(w == 2)ans[i] ++, ans[x] ++;
        else if(w == 1)g[x].pb(i), g[i].pb(x);
    }
    for(int i = 1;i <= N; ++ i)
        for(auto j : g[i])ans[i] += g[j].size();
    
    for(int i = 1;i <= N; ++ i)printf("%lld\n",ans[i] + 1);
    return 0;
}

思路二

每个点的答案由父节点（爷爷节点）来更新，不需要记录具体有哪些子节点，只需要知道子节点的数量就好。

那么每个点i可以到达以下四种点：

• 爷爷节点（如果存在的话）,由点i来更新
• 兄弟节点（父亲的儿子，如果存在的话）由点i来更新
• 自己（儿子的父亲，就算没有儿子根据题意也可以到达）直接 + 1就行，如果有父节点就可以到达父节点
• 孙子节点（如果存在的话）由孙子来更新，所以点i不用管这个

那么现在只剩下两个部分需要更新的，就是点i的爷爷和兄弟。

1.通过fa[fa[i]]来验证爷爷是否存在，若存在那么i和爷爷的答案都 + 1。

2.如果fa[i]存在，那么i的答案还要加上nson[fa[i]] - 1，也就是兄弟节点的个数。注意这里的 - 1不能和前面自身的+ 1抵消，因为父亲未必存在，但是自身不管有没有父亲都是要加上的。

代码二（244ms）

#include <bits/stdc++.h>
#define pb push_back
#define int long long
using namespace std;

inline int readInt()
{
    int s = 0, w = 1;char ch = getchar();
    for(;ch > '9' || ch < '0';ch = getchar())if(ch == '-')w = -1;
    for(;'0' <= ch && ch <= '9';ch = getchar())s = s * 10 + ch - '0';
    return s * w;
}

const int maxn = 1e6 + 10;
int ans[maxn];
int fa[maxn], nson[maxn];


signed main()
{
    ios::sync_with_stdio(0);
	int N = readInt();
    for(int i = 2;i <= N; ++ i)
    {
        int x = readInt(), w = readInt();
        if(w == 2)ans[i] ++, ans[x] ++;
        else if(w == 1)fa[i] = x, nson[x] ++;
    }
    for(int i = 1;i <= N; ++ i)
    {
        if(fa[i])ans[i] += nson[fa[i]], ans[fa[i]] ++;
        if(fa[fa[i]])ans[i] ++, ans[fa[fa[i]]] ++;
    }
    for(int i = 1;i <= N; ++ i)printf("%lld\n",ans[i] + 1);
    return 0;
}