【題解】AtCoder Hitachi 2020 C – ThREE

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

int n, a, b, color[200005];
vector <int> v[200005], colorv[2];

int main(){
    cin >> n;
    for (int i = 1; i < n; i++){
        cin >> a >> b;
        v[a].push_back(b);
        v[b].push_back(a);
    }
    queue <pair<int, int> > q;
    pair<int, int> now;
    q.push({1, 0});
    colorv[0].push_back(1);
    color[1] = 0;
    while (!q.empty()){
        now = q.front();
        q.pop();
        if (color[now.first] == 0){
            for (int i:v[now.first]){
                if (i == now.second) continue;
                color[i] = 1;
                colorv[1].push_back(i);
                q.push({i, now.first});
            }
        }
        else {
            for (int i:v[now.first]){
                if (i == now.second) continue;
                color[i] = 0;
                colorv[0].push_back(i);
                q.push({i, now.first});
            }
        }
    }
    int sz0 = colorv[0].size(), sz1 = colorv[1].size();
    if (sz0 <= n/3 || sz1 <= n/3){
        if (sz0 <= n/3){
            for (int i = 0; i < sz0; i++){
                color[colorv[0][i]] = (i+1) * 3;
            }
            for (int i = (n/3)-sz0-1; i >= 0; i--){
                color[colorv[1][i]] = (sz0+i+1) * 3;
            }
            int num = 1;
            for (int i = (n/3)-sz0; i < sz1; i++){
                if (num % 3 == 0) num++;
                color[colorv[1][i]] = num;
                num++;
            }
        }
        else{
            for (int i = 0; i < sz1; i++){
                color[colorv[1][i]] = (i+1) * 3;
            }
            for (int i = (n/3)-sz1-1; i >= 0; i--){
                color[colorv[0][i]] = (sz1+i+1) * 3;
            }
            int num = 1;
            for (int i = (n/3)-sz1; i < sz0; i++){
                if (num % 3 == 0) num++;
                color[colorv[0][i]] = num;
                num++;
            }
        }
    }
    else{
        int num = 0;
        int cnt1 = n/3, cnt2 = n/3, cnt3 = n/3;
        if (n % 3 == 2) cnt2++;
        if (n % 3 >= 1) cnt1++;
        for (int i = cnt1-1; i >= 0; i--){
            color[colorv[0][i]] = i * 3 + 1;
        }
        for (int i = cnt1; i < sz0; i++){
            color[colorv[0][i]] = cnt3 * 3;
            cnt3--;
        }
        for (int i = cnt2-1; i >= 0; i--){
            color[colorv[1][i]] = i * 3 + 2;
        }
        for (int i = cnt2; i < sz1; i++){
            color[colorv[1][i]] = cnt3 * 3;
            cnt3--;
        }
    }
    for (int i = 1; i <= n; i++){
        cout << color[i] << " ";
    }
    cout << "\n";
}