Codeforces Round #766 (Div. 2) C. Not Assigning

Solution

한 노드에 엣지가 3개 이상이면 prime tree를 생성할 수 없다. 따라서 그래프는 한 줄로 표현이 되야 한다.

#include <iostream>
#include <vector>
#include <map>
#include <queue>
#include <set>

using namespace std;
void dfs(vector<int>& res, int w, int selected, int prev, vector<vector<pair<int, int>>>& edges) {
    int edge = -1, ns = -1;
    for(auto e : edges[selected]) {
        if(e.second != prev) {
            edge = e.first;
            ns = e.second;
            break;
        }
    }
    if(edge == -1) return;
    res[edge] = w == 2 ? 3 : 2;
    dfs(res, res[edge], ns, selected, edges);
}

int main() {
    int rep, n, n1, n2;
    cin>>rep;
    while(rep--) {
        cin>>n;
        vector<vector<pair<int, int>>> edge(n + 1);
        bool overThree = false;
        int single;
        for(int i = 0; i < n - 1; i++) {
            cin>>n1>>n2;
            edge[n1].push_back({i, n2});
            edge[n2].push_back({i, n1});
            overThree |= edge[n1].size() >= 3;
            overThree |= edge[n2].size() >= 3;
        }
        if(overThree) {cout<<"-1"<<endl;continue;}
        for(int i = 1; i <= n; i++) {
            if(edge[i].size() == 1) {
                single = i;
                break;
            }
        }
        vector<int> res(n - 1);
        dfs(res, 2, single, -1, edge);
        for(auto it: res) cout<<it<<" ";cout<<endl;

    }
    return 0;
}