RikkawithGenerals

Today, Rikka is playing a strategic game. She has finished the first stage of the game: She has already established her own country.
Rikka owns n cities, connected by n-1 bidirectional roads. Any two cities are reachable through the roads. Rikka decides to award these cities to n loyal generals. These generals are labeled from 1 to n according to the increasing order of their contributions. Initially, Rikka decides to award the i-th city to the $p_i$ -th general, where $p_1, p_2, \dots, p_n$ is a permutation of length n.
Then, Rikka summons all the generals and shows them the initial plan. Generals are allowed to exchange their cities under two restrictions. General u with city a can change his/her city with general v with city b if and only if:

The contributions of general u and general v are close, i.e. |u-v| should be equal to 1;
The geometric positions of city a and city b are close, i.e. there should be a road between city a and city b.

During the exchange process, one general is allowed to change his/her city many times, and also, one city may be exchanged among many generals.
Not surprised, a quarrel broke out between the generals. It seems that it will take a long time to determine the ownership of the cities. All these things make Rikka bored. To make fun, Rikka wants you to calculate the number of possible award plans.

The first line contains a single integer $t\ (1 \leq t \leq 2 \times 10^5)$ , representing the number of test cases.
For each test case, the first line contains a single integer $n\ (1 \leq n \leq 2\times 10^5)$ , representing the number of cities.
Then n-1 lines follow, each line with two integers $u,v\ (1 \leq u,v \leq n)$ , representing a road between city u and city v.
The last line contains n integers $p_i\ (1 \leq p_i \leq n)$ , representing the initial award plan.
The input guarantees that $p_1, p_2, \dots, p_n$ is a permutation of length n, and $\sum n \leq 2 \times 10^5$ .

#include<bits/stdc++.h> #define fi first #define se second using namespace std; const int N=2e5+100, P=998244353; int T, n, a[N], p[N], las[N], vis[N]; vector< int > nxt[N]; map<int, int > f[N]; void upd(int &x) {x += (x >> 31 & P);} int main() { cin >> T; while(T--) { cin >> n; for(int i=1;i < n;++i) { int x, y; cin >> x >> y; nxt[x].push_back(y); nxt[y].push_back(x); } for(int i=1;i <= n;++i) cin >> p[i], a[p[i]]=i; for(int i=1;i <= n;++i) { for(auto t:nxt[i]) vis[p[t]]=1; las[i]=p[i]; while(vis[las[i]-1]) --las[i]; for(auto t:nxt[i]) vis[p[t]]=0; for(auto &t:nxt[i]) t=p[t]; sort(nxt[i].begin(), nxt[i].end()); } f[1][a[1]]=1; for(int i=1;i < n;++i) for(auto it:f[i]) { int x=it.fi, val=it.se; upd(f[i+1][a[i+1]] += val-P); auto t=upper_bound(nxt[x].begin(), nxt[x].end(), i); if(t != nxt[x].end()) { if(*t == i+1) upd(f[i+1][x] += val-P); else if(las[a[*t]] <= i+1) upd(f[*t][a[*t-1]] += val-P); } } int ans=0; for(auto it:f[n]) upd(ans += it.se-P); printf("%d\n", ans); for(int i=1;i <= n;++i) nxt[i].clear(), f[i].clear(); } return 0; }

详情