USACO Silver 2021 January - Dance Mooves

Official Analysis (C++)

Solution (DSU)

Just like the official analysis, the key observation to make here is that all cows will share the same answer if they have been in the same position. In other words, any cow i will share the same positions with P[i], the position of the cow after K swaps. Since we are only keeping track of the unique items, it suffices to use a set or in this case an unordered_set to keep track of the distinct position a cow covers. After linking all our components together, it suffices to insert them all into a singular set - the set takes care of counting distinct numbers. The trick here is that we perform insert() on components[dsu.get(i)] in order to count this cumulative distinctness in a single position.

Time Complexity: $\mathcal{O}(N\alpha(N) + K)$

#include <bits/stdc++.h> // see /general/running-code-locally
using namespace std;

using ll = long long;

using vi = vector<int>;
#define pb push_back
#define all(x) begin(x), end(x)
#define sz(x) (int)(x).size()