【CodeForces】【搜索】【思維】1218H Function Composition

CodeForces 1218H Function Composition

題目大意

給定一個長度爲$N$的序列$A(1\le A_i\le N)$，現在定義一個函數如下：

• $f(i,1)=A_i$；
• $f(i,m)=A_{f(i,m-1)}$。

給定$m,y$，求有多少個$x$使得$f(x,m)=y$，多組詢問。

分析

考慮將這個問題轉化爲一個圖論問題：當$A_x=y$時，我們連一條從$x$到$y$的有向邊。

對於這個圖，我們仔細觀察一下可以發現如下性質：

• 每個節點的出度都是$1$；
• 每個連通塊都可以被單獨考慮計算貢獻；
• 查詢可以變成找出節點$y$的$m$級祖先有多少個；
• 每個連通塊中只有一個環。

考慮離線處理詢問。

我們先預處理出所有的環。然後枚舉每個位置上的詢問，分兩種情況討論：

• $y$不屬於任何一個環：我們需要統計的就是距離當前深度等於$m$的點的數量，這個可以用 DFS 來完成：先刪掉環上的所有邊，然後在記錄每個深度上的節點數量，就可以求出距離$y$的爲$m$的點。
• $y$在環上，我們將從環上的任意一個節點$z$開始遍歷。設$z$到$y$的距離爲$t$，環長爲$l$。我們先刪掉$z$的出邊（它只有這麼一條），那麼我們就可以比較輕鬆回答詢問了，這樣我們只需要計算距離$y$節點爲$(m+t)\mod l$的節點數量了。

這東西說着容易，寫起來賊難受。。。

參考代碼

#include <queue>
#include <cstdio>
#include <vector>
#include <algorithm>
using namespace std;

typedef long long ll;
const int Maxn = 2e5;

struct DSU {
	int fa[Maxn + 5];
	void init(int n) {
		for(int i = 1; i <= n; i++)
			fa[i] = i;
	}
	int find(int u) {
		return u == fa[u] ? u : fa[u] = find(fa[u]);
	}
	bool unite(int u, int v) {
		u = find(u), v = find(v);
		if(u == v) return false;
		fa[u] = v;
		return true;
	}
};

struct Query {
	int id, typ;
	ll m;
	Query(){}
	Query(int _typ, int _id, ll _m) {
		typ = _typ, id = _id, m = _m;
	}
	bool operator < (const Query &rhs) const {
		return m == rhs.m ? typ > rhs.typ : m > rhs.m;
	}
};

struct Edge {
	int to;
	Edge *nxt;
};

int N, A[Maxn + 5];
DSU s;

Edge pool[Maxn * 2 + 5];
Edge *G[Maxn + 5], *ecnt = &pool[0];
inline void addedge(int u, int v) {
	Edge *p = ++ecnt;
	p->to = v, p->nxt = G[u];
	G[u] = p;
}

vector<int> circles[Maxn + 5];
int cnt_circle;

bool on_circle[Maxn + 5];
bool DFS(int u, int tp) {
	if(u == tp) {
		on_circle[u] = true;
		circles[cnt_circle].push_back(u);
		return true;
	}
	for(Edge *p = G[u]; p != NULL; p = p->nxt) {
		int v = p->to;
		if(!DFS(v, tp)) continue;
		on_circle[u] = true, circles[cnt_circle].push_back(u);
		return true;
	}
	return false;
}

vector<Query> Q[Maxn + 5];

int len, cnt[Maxn + 5];
priority_queue<Query> q;
int mxdep;
int ans[Maxn + 5], res[Maxn + 5];
inline int nxt(int x) {return x == len - 1 ? 0 : x + 1;}
void DFS(int u, int dep, int cirid, int pos) {
	q.push(Query(-1, circles[cirid][pos], dep));
	mxdep = max(mxdep, dep);
	if(!on_circle[u])
		for(int i = 0; i < (int)Q[u].size(); i++)
			if(dep + Q[u][i].m <= Maxn)
				ans[Q[u][i].id] -= cnt[dep + Q[u][i].m];
	cnt[dep]++;
	for(Edge *p = G[u]; p != NULL; p = p->nxt) {
		int v = p->to;
		if(on_circle[v]) continue;
		DFS(v, dep + 1, cirid, nxt(pos));
	}
	if(!on_circle[u])
		for(int i = 0; i < (int)Q[u].size(); i++)
			if(dep + Q[u][i].m <= Maxn)
				ans[Q[u][i].id] += cnt[dep + Q[u][i].m];
}

int main() {
#ifdef LOACL
	freopen("in.txt", "r", stdin);
	freopen("out.txt", "w", stdout);
#endif
	scanf("%d", &N);
	s.init(N);
	for(int i = 1; i <= N; i++) {
		scanf("%d", &A[i]);
		if(!s.unite(A[i], i)) {
			cnt_circle++, DFS(i, A[i]);
			reverse(circles[cnt_circle].begin(), circles[cnt_circle].end());
		}
		addedge(A[i], i);
	}
	int _;
	scanf("%d", &_);
	for(int i = 1; i <= _; i++) {
		ll m;
		int y;
		scanf("%lld %d", &m, &y);
		Q[y].push_back(Query(0, i, m));
	}
	for(int i = 1; i <= cnt_circle; i++) {
		len = (int)circles[i].size();
		for(int j = 0; j < len; j++) {
			int u = circles[i][j];
			mxdep = 0;
			DFS(u, 0, i, j);
			for(int k = 0; k <= mxdep; k++) cnt[k] = 0;
			for(int k = 0; k < (int)Q[u].size(); k++) {
				int v = circles[i][(j + Q[u][k].m) % len];
				q.push(Query(Q[u][k].id, v, Q[u][k].m));
			}
		}
		while(!q.empty()) {
			Query tmp = q.top();
			q.pop();
			if(tmp.typ == -1) res[tmp.id]++;
			else ans[tmp.typ] = res[tmp.id];
		}
	}
	for(int i = 1; i <= _; i++)
		printf("%d\n", ans[i]);
	return 0;
}