cf1787 D. Game on Axis(思维+建图)

传送门

#include <bits/stdc++.h>

using namespace std;

#define endl "\n"
#define debug cout<<"debug"<<endl
typedef long long ll;
typedef long double ld;
typedef pair<int, int> PII;
const double eps = 1e-8;
const double PI = 3.14159265358979323;
const int N = 2e5+10, M = 2*N, mod = 998244353;
const int INF = 0x3f3f3f3f;

int n;
int a[N];
int sz[N], fa[N];
vector<int> v[N];

void dfs(int u)
{
    sz[u] = 1;
    for(auto t : v[u])
    {
        dfs(t);
        sz[u] += sz[t];
    }
}

void solve()
{
    cin>>n;
    for(int i = 0; i<=n; i++) sz[i] = 0, v[i].clear();
    // 将路径建图, 终点设为0, 不难发现只有以0为根的树才能到达终点, 其他点一定在某个基环树上
    // 要使1能到达终点, 分两种情况
    // 若1的根为0, 则1一定能到达终点, 不在1-->0路径的节点可以任意修改不影响, 只需要考虑在路径上的点
    // 而在路径上的点只要修改到它的子树上就会出现基环导致不合法, 所以可以枚举路径上的点计算不合法数(res = 总数-不合法数)即可
    // 若1不在以0为根的点, 那么1肯定在某个基环树上, 则只能修改1本身 或 环中1可达的点, 且只能修改到以0为根的这棵树上
    for(int i = 1; i<=n; i++)
    {
        cin>>a[i];
        int p = i+a[i];
        if(p>n || p<=0) p = 0;
        v[p].push_back(i);
        fa[i] = p;
    }
    // dfs 0为根的树, 求子树大小, 若1在树上, 则sz[1]>=1, 否则sz[1]=0;
    dfs(0);
    if(sz[1])
    {
        ll res = (ll)n*(2ll*(ll)n+1);
        int u = 1;
        while(u)
        {
            // 不合法数, 即自身子树节点或其他基环树上的点
            res -= sz[u] + (n - (sz[0] - 1));
            u = fa[u];
        }
        cout<<res<<endl;
    }else{
        int st[n+1];
        for(int i = 0; i<=n; i++) st[i] = 0;
        ll res = 0;
        int u = 1;
        while(!st[u])
        {
            st[u] = 1;
            // 修改范围为0为根的子树的点
            res += sz[0]-1+n+1;
            u = fa[u];
        }
        cout<<res<<endl;
    }
}

int main()
{
    ios::sync_with_stdio(false),cin.tie(0),cout.tie(0);
    int T;
    T = 1;
    cin>>T;
    while(T -- )
    {
        solve();
    }
    return 0;
}