[ICPC2019 WF]Hobson's Trains

传送


这道题虽然思维难度不高，但是我却收获颇多。


这个图的模型就是好多棵基环内向树，然后求对于每一个点\(u\)，在\(k\)步之内能走到\(u\)的点的数量（包括点\(u\)自己）。


首先自然是分成两部分：树上的点和环上的点。
对于树上的每一个点，无非是把他的的\(k\)级之内的祖先都加1，这个用树状数组或树上差分都能解决；
对于环上的每一个点，要判断能否绕完一圈，如果能，就加上环的长度；否则于把每一个点走\(k\)步之内的点的值都加1，这个用差分也很好解决。
所以这道题就做完了。


这是常规的思路，但却不是最优化、最简洁的思路。对于树上的点来说，无论是树状数组还是树上差分都是\(O(nlogn)\)的（树上差分需要预处理倍增数组），但实际上能做到\(O(n)\)，而且额外简单。
我们只要在树上dfs的时候维护一个栈，但这个栈中的每一个元素都有用，st[i]表示深度为\(i\)的点的个数。
具体步骤是在第一次dfs到点\(u\)的时候记录栈中的所有元素个数\(num_1\)，在回溯到\(u\)的时候在记录一次\(num_2\)，期间把深度大于\(d_u+k\)的点都踢出栈，这样\(num_2 - num_1\)就是距离\(u\)小于等于\(k\)的点的个数了！


处理完树上的点后，还有一些点延伸到环上，即栈中剩下的点，那么只要仿照环上的点进行差分就可以了。

#include<cstdio>
#include<iostream>
#include<cmath>
#include<algorithm>
#include<cstring>
#include<cstdlib>
#include<cctype>
#include<vector>
#include<queue>
#include<assert.h>
#include<ctime>
using namespace std;
#define enter puts("") 
#define space putchar(' ')
#define Mem(a, x) memset(a, x, sizeof(a))
#define In inline
#define forE(i, x, y) for(int i = head[x], y; ~i && (y = e[i].to); i = e[i].nxt)
typedef long long ll;
typedef double db;
const int INF = 0x3f3f3f3f;
const db eps = 1e-8;
const int maxn = 5e5 + 5;
In ll read()
{
	ll ans = 0;
	char ch = getchar(), las = ' ';
	while(!isdigit(ch)) las = ch, ch = getchar();
	while(isdigit(ch)) ans = (ans << 1) + (ans << 3) + ch - '0', ch = getchar();
	if(las == '-') ans = -ans;
	return ans;
}
In void write(ll x)
{
	if(x < 0) x = -x, putchar('-');
	if(x >= 10) write(x / 10);
	putchar(x % 10 + '0');
}
In void MYFILE()
{
#ifndef mrclr
	freopen("ha2.in", "r", stdin);
	freopen("ha.out", "w", stdout);
#endif
}

int n, K, To[maxn], ans[maxn];
struct Edge
{
	int nxt, to;
}e[maxn];
int head[maxn], ecnt = -1;
In void addEdge(int x, int y)
{
	e[++ecnt] = (Edge){head[x], y};
	head[x] = ecnt;
}

int st[maxn], top = 0;
In void dfs(int& num, int now, int _f, int d)
{
	int tp = num;
	if(d > top) st[++top] = 0;
	num++, st[d]++;
	forE(i, now, v) if(v ^ _f) dfs(num, v, now, d + 1);
	ans[now] = num - tp;
	if(top == d + K) num -= st[top--];		//对于u的父亲来说，这些点的距离已经超过k了 
}

int a[maxn], dif[maxn], acnt = 0;
In int nxt(int i, int x) {return i + x <= acnt ? i + x : i + x - acnt;}
In void solve(int x)
{
	int u = x, v = To[x];
	while(u != v && u != To[v]) u = To[u], v = To[To[v]];	//用floyd找环法！一行就搞定了！ 
	a[acnt = 1] = u;
	for(int v = To[u]; v != u; v = To[v]) a[++acnt] = v;
	fill(dif + 1, dif + acnt + 1, 0);
	int sum = min(K + 1, acnt);
	for(int i = 1; i <= acnt; ++i)
		forE(k, a[i], v)
		{
			if(v == a[nxt(i, acnt - 1)]) continue;
			int num = 0;
			st[top = 0] = 0, dfs(num, v, a[i], 1);
			for(int d = 1; d <= top; ++d)	          //对于树上深度小于等于k的点要考虑对环的贡献 
			{
				int dnum = st[d];
				if(K - d >= acnt - 1) sum += dnum;//如果已经能走到环上每一个点了 
				else
				{
					int j = nxt(i, K - d);	  //否则进行环上差分：如果绕到了i的前面，把dif[1]+=dnum 
					dif[i] += dnum, dif[nxt(j, 1)] -= dnum;
					if(nxt(j, 1) < i) dif[1] += dnum;
				}
			}
		}
	for(int i = 1; i <= acnt; ++i) sum += dif[i], ans[a[i]] = sum;
}

int main()
{
//	MYFILE();
	Mem(head, -1), ecnt = -1;
	n = read(), K = read();
	for(int i = 1; i <= n; ++i)
	{
		int x = read();
		addEdge(x, i), To[i] = x;
	}
	fill(ans + 1, ans + n + 1, -1);
	for(int i = 1; i <= n; ++i) if(ans[i] == -1) solve(i);
	for(int i = 1; i <= n; ++i) write(ans[i]), enter;
	return 0;
}