Gym - 101615 D Rainbow Roads dfs序

题目传送门

题目大意：

      给出一颗树，每条边都有一个颜色，对一个点来说，如果其他所有点到这个点的简单路径，相连的边颜色都不同，这个点即合法点，统计所有的合法点。

思路：

      对于一个节点来说

      1、如果这个节点的两个子节点的边颜色一样，那么这两个子节点的子树都要放弃。

      2、如果这个节点的子节点和他的父节点的边的颜色一样，那么子节点所在子树放弃，父节点中除了这个节点以外的其他节点都要放弃。

剩下的点就是合法点。

      那怎么做到放弃呢？

      先dfs处理出每个节点的dfs序，每个dfs序对应的节点标号，每个节点的子树大小（包括本身），对于第一种情况来说，dfs序在  子树节点dfs序  到 子树节点dfs序加上子树大小  这个左闭右开的区间都要放弃，则用一个数组，左边界加一，右边界减一。

对于第二种情况，放弃的子节点和上述一样处理，父节点就要放弃  dfs序从1到父节点dfs序，父节点dfs序+1到最后所有的点，则对应位置加一或者减一，最后vis[ i ]=vis[ i  -1 ] + arr[ i ].处理完后，vis还等于0的点，就是合法的dfs序，找到这个dfs序对应的节点就可以了（dfs的时候就记录过了，大佬说这个做法叫差分约束）。

这题五万个点，每个点三个数据，用读入优化居然没有快，，，可能是我的读入优化还不够优秀吧。

#include<iostream>
#include<cstdio>
#include<cmath>
#include<algorithm>
#include<string.h>
#include<sstream>
#include<set>
#include<map>
#include<vector>
#include<queue>
#include<stack>
#include<bitset>
#define CLR(a,b) memset((a),(b),sizeof((a))) 
using namespace std;
typedef long long ll;
inline int rd() {
	int f = 1; int x = 0; char s = getchar();
	while (s<'0' || s>'9') { if (s == '-')f = -1; s = getchar(); }
	while (s >= '0'&&s <= '9') { x = x * 10 + s - '0'; s = getchar(); }x *= f;
	return x;
}
int n;
struct edge {
	int v, color;
	edge(){}
	edge(int v,int color): v(v),color(color){}
	
};
bool cmp(edge &a, edge &b)
{
	return a.color < b.color;
}
const int maxn = 50010;
vector<edge >g[maxn];
int arr[maxn], fa[maxn], son[maxn], dfn[maxn],tot,cnt,vis[maxn],dir[maxn];
void dfs(int u, int pre)
{
	fa[u] = pre;
	son[u] = 1;
	dfn[u] = ++cnt;
	dir[cnt] = u;
	for (auto it : g[u])
	{
		if (it.v == pre)continue;
		dfs(it.v, u);
		son[u] += son[it.v];
	}
}
int main() {
	cin >> n;
	for (int i = 1; i < n; i++)
	{
		int u, v, color;
		u = rd(), v = rd(), color = rd();
		//scanf("%d%d%d", &u, &v, &color);
		g[u].push_back(edge(v, color));
		g[v].push_back(edge(u, color));
	}
	dfs(1, 0);
	for (int i = 1; i <= n; i++)
	{
		sort(g[i].begin(), g[i].end(), cmp);
		int si = g[i].size();
		for (int j = 0; j < si - 1; j++)
		{
			if (g[i][j].color == g[i][j + 1].color)
			{
				
				int x = g[i][j].v, y = g[i][j+1].v;
				if (fa[x] == i && fa[y] == i)
				{
					
					arr[dfn[x]]++;
					arr[dfn[x] + son[x]]--;
					arr[dfn[y]]++;
					arr[dfn[y] + son[y]]--;
				}
				else if (fa[x] == i && fa[i] == y)
				{
					
					arr[dfn[x]]++;
					arr[dfn[x] + son[x]]--;
					arr[dfn[1]]++;
					arr[dfn[i]]--;
					arr[dfn[i] + son[i]]++;
				}
				else if (fa[y] == i && fa[i] == x)
				{
					swap(x, y);
					arr[dfn[x]]++;
					arr[dfn[x] + son[x]]--;
					arr[dfn[1]]++;
					arr[dfn[i]]--;
					arr[dfn[i] + son[i]]++;
				}
			}
		}
	}
	for (int i = 1; i <= n; i++)
	{
		vis[i] = vis[i - 1] + arr[i];
	}
	vector<int >ans;
	for (int i = 1; i <= n; i++)
	{
		if (vis[i] == 0)
		{
			tot++;
			ans.push_back(dir[i]);
		}
	}
	printf("%d\n", tot);
	sort(ans.begin(), ans.end());
	for (auto it : ans)
	{
		printf("%d\n", it);
	}
}