제출 #916874

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
916874		GrindMachine	Unique Cities (JOI19_ho_t5)	C++17	100 / 100	587 ms	100564 KiB

joi2019_ho_t5

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

using namespace std;
using namespace __gnu_pbds;

template<typename T> using Tree = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;
typedef long long int ll;
typedef long double ld;
typedef pair<int,int> pii;
typedef pair<ll,ll> pll;

#define fastio ios_base::sync_with_stdio(false); cin.tie(NULL)
#define pb push_back
#define endl '\n'
#define sz(a) (int)a.size()
#define setbits(x) __builtin_popcountll(x)
#define ff first
#define ss second
#define conts continue
#define ceil2(x,y) ((x+y-1)/(y))
#define all(a) a.begin(), a.end()
#define rall(a) a.rbegin(), a.rend()
#define yes cout << "Yes" << endl
#define no cout << "No" << endl

#define rep(i,n) for(int i = 0; i < n; ++i)
#define rep1(i,n) for(int i = 1; i <= n; ++i)
#define rev(i,s,e) for(int i = s; i >= e; --i)
#define trav(i,a) for(auto &i : a)

template<typename T>
void amin(T &a, T b) {
    a = min(a,b);
}

template<typename T>
void amax(T &a, T b) {
    a = max(a,b);
}

#ifdef LOCAL
#include "debug.h"
#else
#define debug(x) 42
#endif

/*

refs:
edi
https://codeforces.com/blog/entry/65042?#comment-491880

key idea:
for a node u, let d = arbitrary node with farthest dis from u
all good nodes lie on the path from u to d

proof:
consider the tree rooted at the path (u,d)
anybody in a branching subtree cant be good
why? because there exists another node on the path (u,d) with the same distance
if there exists a good guy in the subtree, then it contradicts the definition of d (d should have been somewhere in the subtree)

(s,t) = endpoints of diameter
farthest node from u = s or t

run dfs from s and t, calculate answers

when we go into a subtree, find the deepest in the other subtree and remove the respective nodes from consideration 
(look at the pictures in the edi for better understanding)

in order to find #of good nodes, maintain a lazy segtree with range add updates and min, no.of min queries
when going into a subtree, add 1 to all bad nodes that are caused by the deepest in the other subtree
we need #of 0s
we know that min_val >= 0, so just get min and #of min

how to count #of distinct guys?

similar idea to ceoi harbingers
maintain a stack of all the "unique nodes" (i.e nodes with colors that are not present before its position in the stack) when doing the dfs
when we want to add a new node, add him only if the color is not present in the stack
keep an additional variable called "siz" which maintains the size of the active part of the stack
when we want to remove multiple nodes, just change the value of siz
rollback the changes after dfs

*/

const int MOD = 1e9 + 7;
const int N = 2e5 + 5;
const int inf1 = int(1e9) + 5;
const ll inf2 = ll(1e18) + 5;

vector<ll> adj[N];
vector<ll> a(N);
pll diam = {-1,-1};

void dfs1(ll u, ll p, ll d){
    pll px = {d,u};
    amax(diam,px);
    trav(v,adj[u]){
        if(v == p) conts;
        dfs1(v,u,d+1);
    }
}

ll dis[N][2];
ll deepest[N][2];

void dfs2(ll u, ll p, ll t){
    trav(v,adj[u]){
        if(v == p) conts;
        dis[v][t] = dis[u][t]+1;
        dfs2(v,u,t);
        amax(deepest[u][t],deepest[v][t]+1);
    }
}

vector<ll> ans(N);
vector<ll> mx_without(N);
vector<ll> first_occ(N,inf2);
vector<pll> active(N);
ll siz = 0;

void dfs3(ll u, ll p, ll d, ll t){
    ll mx = deepest[u][t];
    ll cnt = upper_bound(active.begin(),active.begin()+siz,make_pair(d-mx,-1ll))-active.begin();
    if(dis[u][t] >= dis[u][t^1]){
        ans[u] = cnt;
    }

    mx = 0;
 
    trav(v,adj[u]){
        if(v == p) conts;
        mx_without[v] = mx;
        amax(mx,deepest[v][t]+1);
    }
 
    mx = 0;
    reverse(all(adj[u]));
 
    trav(v,adj[u]){
        if(v == p) conts;
        amax(mx_without[v],mx);
        amax(mx,deepest[v][t]+1);
    }
 
    reverse(all(adj[u]));
    
    trav(v,adj[u]){
        if(v == p) conts;

        vector<array<ll,4>> history;

        mx = mx_without[v];
        ll new_siz = upper_bound(active.begin(),active.begin()+siz,make_pair(d-mx,-1ll))-active.begin();
        ll old_siz = siz;
        history.pb({3,siz,-1,-1});
        siz = new_siz;

        ll x = a[u];

        if(first_occ[x] >= siz){
            history.pb({1,siz,active[siz].ff,active[siz].ss});
            history.pb({2,x,first_occ[x],-1});
            history.pb({3,siz,-1,-1});

            ll w = active[siz].ss;
            if(first_occ[a[w]] == siz){            
                history.pb({2,a[w],first_occ[a[w]],-1});
                first_occ[a[w]] = inf2;
            }

            active[siz] = {d,u};
            first_occ[x] = siz;
            siz++;
        }

        dfs3(v,u,d+1,t);

        reverse(all(history));

        trav(ar,history){
            auto [t,ind,val1,val2] = ar;
            if(t == 1){
                active[ind] = {val1,val2};
            }
            else if(t == 2){
                first_occ[ind] = val1;
            }
            else{
                siz = ind;
            }
        }
    }
}

void solve(int test_case)
{
    ll n,m; cin >> n >> m;
    rep1(i,n-1){
        ll u,v; cin >> u >> v;
        adj[u].pb(v), adj[v].pb(u);
    }
    rep1(i,n) cin >> a[i];

    dfs1(1,-1,0);
    ll s = diam.ss;
    diam = {-1,-1};
    dfs1(s,-1,0);
    ll t = diam.ss;

    dfs2(s,-1,0);
    dfs2(t,-1,1);

    dfs3(s,-1,0,0);
    dfs3(t,-1,0,1);

    rep1(i,n){
        cout << ans[i] << endl;
    }
}

int main()
{
    fastio;

    int t = 1;
    // cin >> t;

    rep1(i, t) {
        solve(i);
    }

    return 0;
}

컴파일 시 표준 에러 (stderr) 메시지

joi2019_ho_t5.cpp: In function 'void dfs3(ll, ll, ll, ll)':
joi2019_ho_t5.cpp:158:12: warning: unused variable 'old_siz' [-Wunused-variable]
  158 |         ll old_siz = siz;
      |            ^~~~~~~

• Subtask 1: 4 / 4
• Subtask 2: 32 / 32
• Subtask 3: 32 / 32
• Subtask 4: 32 / 32