Submission #1174558

#	Time	Username	Problem	Language	Result	Execution time	Memory
1174558		Pacybwoah	Balanced Tree (info1cup18_balancedtree)	C++20	13 / 100	74 ms	7612 KiB

balancedtree

#include<iostream>
#include<vector>
#include<algorithm>
#include<utility>
#include<set>
#include<cmath>
using namespace std;
typedef long long ll;
const int inf = 1e9;
struct reroot{
    vector<vector<int>> graph;
    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 add(int a, int b){
        graph[a].push_back(b);
        graph[b].push_back(a);
    }
    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){
        multiset<int> s;
        s.insert(up + 1);
        for(auto &x: graph[node]){
            if(x == parent) continue;
            s.insert(dp1[x] + 1);
        }
        dp2[node] = (*s.begin());
        if(vec[node]) s.insert(0);
        for(auto &x: graph[node]){
            if(x == parent) continue;
            s.erase(s.find(dp1[x] + 1));
            dfs2(x, node, (*s.begin()));
            s.insert(dp1[x] + 1);
        }
    }
    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;
        reroot zer(n), one(n);
        for(int i = 1; i < n; i++){
            int a, b;
            cin >> a >> b;
            zer.add(a, b);
            one.add(a, b);
        }
        vector<int> vec(n + 1);
        for(int i = 1; i <= n; i++){
            cin >> vec[i];
            if(vec[i] == 1) one.set_1(i);
            else zer.set_1(i);
        }
        zer.do_dp();
        one.do_dp();
        int ans = 0;
        for(int i = 1; i <= n; i++){
            if(vec[i] == 1){
                ans = max(ans, one.dp2[i]);
            }
            else{
                ans = max(ans, zer.dp2[i]);
            }
        }
        if(ans >= inf) cout << "-1\n";
        else cout << ans << "\n";
        if(ans < inf){
            for(int i = 1; i <= n; i++) cout << vec[i] << ' ';
            cout << "\n";
        }
    }
}
// g++ -std=c++17 -Wall -Wextra -Wshadow -fsanitize=undefined -fsanitize=address -o run pD.cpp -g

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