rt,思路是分成出发点在树或基环树考虑。树的话总起点到根节点路径上的每一个点都不能连回子树或其他基环树,剩下随便;基环树则路径上的每个点都不能连回自身,也不能连向其他基环树,剩下随便连。但是 WA on pretest 2……
求大佬看看
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int MAXN = 2e5 + 10;
struct node {
int v, nxt;
} e[MAXN];
int head[MAXN], tot;
inline
void add(int u, int v) {
e[++tot] = { v, head[u] }, head[u] = tot;
}
bool vis[MAXN];
int size, cnt, a[MAXN], s[MAXN];
void dfs(int u) {
if (vis[u]) return ; vis[u] = 1, size++;
for (int i = head[u]; i; i = e[i].nxt) cnt++, dfs(e[i].v);
}
int t, n, x, y, k, rt;
void dfs2(int u) {
if (vis[u]) return cnt++, void();
if (u <= 0 || u > n) return ;
vis[u] = 1, size++, dfs2(a[u]);
}
void dfs3(int u, int f) {
s[u] = 1;
for (int i = head[u], v; i; i = e[i].nxt) {
v = e[i].v;
if (v != f) dfs3(v, u), s[u] += s[v];
}
}
ll ans;
int main() {
for (scanf("%d", &t); t--;) {
scanf("%d", &n), x = y = 0;
for (int i = 1; i <= n; i++) head[i] = vis[i] = 0; tot = 0;
for (int i = 1; i <= n; i++) {
scanf("%d", &a[i]), a[i] += i;
if (a[i] > 0 && a[i] <= n) add(i, a[i]), add(a[i], i);
}
for (int i = 1; i <= n; i++) {
if (!vis[i]) size = cnt = 0, dfs(i), size == cnt >> 1 ? x += size : y += size;
if (i == 1) k = size;
}
for (int i = 1; i <= n; i++) vis[i] = 0; size = cnt = 0, dfs2(1);
if (cnt) ans = (ll)size * (2 * n + 1 - x) + (ll)(n - size) * (2 * n + 1);
else {
for (rt = 1; a[rt] > 0 && a[rt] <= n; rt = a[rt]); dfs3(rt, 0);
ans = (ll)(n - size) * (2 * n + 1);
for (rt = 1; rt > 0 && rt <= n; rt = a[rt]) ans += 2 * n + 1 - x - s[rt];
}
printf("%lld\n", ans);
}
}