Codeforces Manthan Codefest 18 (Div. 1 + Div. 2) D. Valid BFS? (R1700)

こういうの超苦手。。。解き方はわかるけど実装を工夫する感じの。

問題へのリンク

問題概要

N 頂点のツリー (頂点番号は 1, 2, ..., N) と、(1, 2, ...., N) の順列 a が与えられる。

このツリーを頂点 1 から出発して BFS したとして、その頂点訪問順序が a になることがありうるかどうかを判定せよ。

制約

• 1 <= N <= 2 × 10⁵

考えたこと

訪問順序 a にしたがって BFS をシミュレートすればよい。訪問順序 a にしたがって BFS しようとしたときに、a[i] が果たして親から見てちゃんと子供になっているかどうかを判定していく感じ。

#include <iostream>
#include <vector>
#include <set>
#include <queue>
using namespace std;

int N;
typedef vector<vector<int> > Graph;
typedef pair<int,int> pint;
Graph G;
set<pint> se;
vector<int> deg;
vector<int> a;

bool solve() {
    queue<int> que;
    que.push(0);
    
    int iter = 1;
    if (a[0] != 0) return false;
    while (!que.empty()) {
        int v = que.front(); que.pop();
        for (int i = 0; i < deg[v]; ++i) {
            if (iter >= N) return false;
            int nv = a[iter++];
            if (!se.count(pint(v, nv))) return false;
            que.push(nv);
        }
    }
    
    return true;
}

int main() {
    while (cin >> N) {
        G.clear(); G.resize(N); se.clear(); deg.assign(N, 0);
        for (int i = 0; i < N-1; ++i) {
            int u, v; cin >> u >> v; --u, --v;
            G[u].push_back(v);
            G[v].push_back(u);
            se.insert(pint(u, v));
            se.insert(pint(v, u));
            deg[u]++;
            deg[v]++;
        }
        for (int v = 0; v < N; ++v) {
            sort(G[v].begin(), G[v].end());
        }
        a.resize(N);
        for (int i = 0; i < N; ++i) cin >> a[i], --a[i];
        for (int v = 1; v < N; ++v) --deg[v];
        if (solve()) cout << "Yes" << endl;
        else cout << "No" << endl;
    }
}