Witch game

//#define _GLIBCXX_DEBUG
#include <bits/stdc++.h>
using namespace std;
#define pb push_back
#define mp make_pair
#define fst first
#define snd second
#define forn(i, n) for (int i = 0; i < int(n); ++i)
typedef long long ll;
typedef vector<int> vi;
typedef vector<vi> vvi;
typedef vector<ll> vll;
typedef pair<int, int> pii;
typedef vector<pii> vii;
#define sz(c) (int)(c).size()
#define ALL(c) (c).begin(), (c).end()
 
void solve (int n)
{
    vi p(n);
    vector<set<int>> g(n);
    forn (i, n)
    {
        cin >> p[i];
        --p[i];
        g[i].insert(p[i]);
        g[p[i]].insert(i);
    }
 
    ll deg1 = 0, tris = 0;
    forn (i, n) deg1 += (ll)sz(g[i]) * (ll)(n - 1 - sz(g[i]));
    forn (i, n)
    {
        int j = p[i], k = p[j];
        if (j != i && j != k && i != k && p[k] == i)
            tris++;
    }
 
    assert(deg1 % 2LL == 0 && tris % 3LL == 0);
    ll ans = deg1 / 2LL + tris / 3LL;
    ans = (ll)n * (ll)(n - 1) * (ll)(n - 2) / 6LL - ans;
    cout << ans << '\n';
}
 
int main()
{
    #ifndef ONLINE_JUDGE
    //freopen("input.txt", "rt", stdin);
    //freopen("output.txt", "w", stdout);
    #endif // ONLINE_JUDGE
 
    ios_base::sync_with_stdio(false);
    cin.tie(0);
 
    int n;
    while (cin >> n)
        solve(n);
}

Test 1

Group: 1

Verdict: ACCEPTED

Test 2

Group: 1

Verdict: ACCEPTED

Test 3

Group: 1

Verdict: ACCEPTED

Test 4

Group: 1

Verdict: ACCEPTED

Test 5

Group: 1

Verdict: ACCEPTED

Test 6

Group: 2

Verdict: ACCEPTED

Test 7

Group: 2

Verdict: ACCEPTED

Test 8

Group: 2

Verdict: ACCEPTED

Test 9

Group: 2

Verdict: ACCEPTED

Test 10

Group: 2

Verdict: ACCEPTED

Test 11

Group: 3

Verdict: ACCEPTED

Test 12

Group: 3

Verdict: ACCEPTED

Test 13

Group: 3

Verdict: ACCEPTED

Test 14

Group: 3

Verdict: ACCEPTED

Test 15

Group: 3

Verdict: ACCEPTED