Submission #1174698

#	Time	Username	Problem	Language	Result	Execution time	Memory
1174698		Pacybwoah	Balanced Tree (info1cup18_balancedtree)	C++20	13.50 / 100	38 ms	2156 KiB

balancedtree

#pragma GCC optimize("O3,unroll-loops")
#pragma GCC target("bmi,bmi2,avx,avx2")
#include<iostream>
#include<vector>
#include<algorithm>
#include<utility>
#include<set>
#include<cmath>
using namespace std;
typedef long long ll;
const int inf = 1e9;
vector<vector<int>> graph;
struct reroot{
    vector<int> vec, dp1, dp2;
    int n;
    reroot(int _n){
        n = _n;
        graph.resize(n + 1);
        vec.resize(n + 1);
        dp1.resize(n + 1);
        dp2.resize(n + 1);
    }
    void set_1(int x){
        vec[x] = 1;
    }
    void dfs1(int node, int parent){
        if(vec[node]) dp1[node] = 0;
        else dp1[node] = inf;
        for(auto &x: graph[node]){
            if(x == parent) continue;
            dfs1(x, node);
            dp1[node] = min(dp1[node], dp1[x] + 1);
        }
    }
    void dfs2(int node, int parent, int up){
        dp2[node] = up + 1;
        int sz = (int)graph[node].size();
        if(node != 1) sz--;
        vector<int> pre(sz + 1, inf), suf(sz + 1);
        int cnt = 0;
        for(auto &x: graph[node]){
            if(x == parent) continue;
            dp2[node] = min(dp2[node], dp1[x] + 1);
            cnt++;
            pre[cnt] = dp1[x] + 1;
            suf[cnt] = dp1[x] + 1;
        }
        //cout << node << " " << dp2[node] << endl;
        for(int i = 1; i <= sz; i++) pre[i] = min(pre[i], pre[i - 1]);
        for(int i = sz - 1; i > 0; i--) suf[i] = min(suf[i], suf[i + 1]);
        cnt = 0;
        for(auto &x: graph[node]){
            if(x == parent) continue;
            cnt++;
            int down = up + 1;
            if(vec[node]) down = 0;
            down = min(down, pre[cnt - 1]);
            if(cnt < sz) down = min(down, suf[cnt + 1]);
            dfs2(x, node, down);
        }
    }
    void do_dp(){
        dfs1(1, 1);
        dfs2(1, 1, inf);
    }
};
int main(){
    ios::sync_with_stdio(false);
    cin.tie(0);
    int t;
    cin >> t;
    while(t--){
        int n;
        cin >> n;
        vector<vector<int>>().swap(graph);
        graph.resize(n + 1);
        for(int i = 1; i < n; i++){
            int a, b;
            cin >> a >> b;
            graph[a].push_back(b);
            graph[b].push_back(a);
        }
        vector<int> vec(n + 1);
        for(int i = 1; i <= n; i++) cin >> vec[i];
        reroot zer(n), one(n);
        for(int i = 1; i <= n; i++){
            if(vec[i] == 0){
                zer.set_1(i);
            }
            else if(vec[i] == 1){
                one.set_1(i);
            }
            else{
                zer.set_1(i);
                one.set_1(i);
            }
        }
        int ans = 0;
        zer.do_dp();
        one.do_dp();
        for(int i = 1; i <= n; i++){
            if(vec[i] == 0){
                ans = max(ans, zer.dp2[i]);
            }
            else if(vec[i] == 1){
                ans = max(ans, one.dp2[i]);
            }
        }
        if(ans >= inf){
            cout << "-1\n";
            continue;
        }
        cout << ans + 2 << "\n";
        for(int i = 1; i <= n; i++){
            if(vec[i] == 1) cout << "1 ";
            else cout << "0 ";
        }
        cout << '\n';
    }
}
// g++ -std=c++17 -Wall -Wextra -Wshadow -fsanitize=undefined -fsanitize=address -o run pD.cpp -g

• Subtask 1: 0 / 10
• Subtask 2: 0 / 13
• Subtask 3: 0 / 26
• Subtask 4: 13.5 / 45
• Subtask 5: 0 / 6