問題概要 (意訳)

要素数 $N$ の数列 $A_{1}, A_{2}, \dots, A_{N}$ があって、はじめは全要素が 0 となっている。以下の操作を好きな回数だけ行って、数列を $B_{1}, B_{2}, \dots, B_{N}$ の状態にすることが可能かどうかを判定せよ。

$1, 2, \dots, N$ の中から $K$ 個を選ぶ
これらの要素でサイクルを組み、 $l_{1} \to l_{2} \to \dots \to l_{K} \to l_{1}$ とする
$k = 1, 2, \dots, K$ に対して、 $A_{l_{k}}$ を $l_{k+1}$ に書き換える

制約

$1 \le K \le N \le 2 \times 10^{5}$
$1 \le B_{i} \le N$
マルチテストケースである

考えたこと

数列 $B$ で Functional Graph を作る。このグラフのどの連結成分にもサイクルがちょうど 1 つだけ存在することとなる。

これらのサイクルのサイズがすべて $K$ → Yes
そうでない → No

となる。計算量は $O(N)$ 。

コード

Functional Graph のサイクル列挙のライブラリを整備していたので、それを貼るだけで AC できた！

#include <bits/stdc++.h>
using namespace std;

// Edge Class
template<class T> struct Edge {
    int from, to;
    T val;
    Edge() : from(-1), to(-1), val(-1) { }
    Edge(int f, int t, T v = -1) : from(f), to(t), val(v) {}
    friend ostream& operator << (ostream& s, const Edge& E) {
        return s << E.from << "->" << E.to;
    }
};

// G[v] := é ç¹ v ããåºã¦ããè¾º
template<class T> struct CycleDetection {
    // input
    vector<Edge<T>> G;
    
    // intermediate results
    vector<bool> seen, finished;
    vector<int> history;
    
    // constructor
    CycleDetection() { }
    CycleDetection(const vector<Edge<T>> &graph) { init(graph); }
    void init(const vector<Edge<T>> &graph) {
        G = graph;
        seen.assign(G.size(), false);
        finished.assign(G.size(), false);
    }
    
    // return the vertex where cycle is detected
    int search(int v) {
        do {
            seen[v] = true;
            history.push_back(v);
            v = G[v].to;
            if (finished[v]) {
                v = -1;
                break;
            }
        } while (!seen[v]);
        pop_history();
        return v;
    }
    
    // pop history
    void pop_history() {
        while (!history.empty()) {
            int v = history.back();
            finished[v] = true;
            history.pop_back();
        }
    }
    
    // reconstruct
    vector<Edge<T>> reconstruct(int pos) {
        // reconstruct the cycle
        vector<Edge<T>> cycle;
        int v = pos;
        do {
            cycle.push_back(G[v]);
            v = G[v].to;
        } while (v != pos);
        return cycle;
    }
    
    // find cycle, v is the start vertex
    vector<Edge<T>> detect_from_v(int v) {
        int pos = search(v);
        if (pos != -1) return reconstruct(pos);
        else return vector<Edge<T>>();
    }
    
    // find all cycle
    vector<vector<Edge<T>>> detect_all() {
        vector<vector<Edge<T>>> res;
        for (int v = 0; v < (int)G.size(); ++v) {
            if (finished[v]) continue;
            int pos = search(v);
            if (pos == -1) continue;
            const vector<Edge<T>> &cycle = reconstruct(pos);
            if (!cycle.empty()) res.push_back(cycle);
        }
        return res;
    }
};

void solve() {
    int N, K;
    cin >> N >> K;
    vector<int> A(N);
    for (int i = 0; i < N; ++i) {
        cin >> A[i];
        --A[i];
    }
    
    if (K == 1) {
        for (int i = 0; i < N; ++i) if (A[i] != i) {
            cout << "NO" << endl;
            return;
        }
    }
    
    // Functional Graph æ§ç¯
    vector<Edge<long long>> G(N);
    for (int i = 0; i < N; ++i) G[i] = Edge<long long>(i, A[i], i);
    
    // éè·¯æ¤åº
    using Cycle = vector<Edge<long long>>;
    CycleDetection<long long> cd(G);
    const vector<Cycle> &cycles = cd.detect_all();
    
    // éè¨
    long long res = 0;
    for (const auto &cycle : cycles) {
        if (cycle.size() != K) {
            cout << "NO" << endl;
            return;
        }
    }
    cout << "YES" << endl;
}

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);

    int T;
    cin >> T;
    while (T--) solve();
}