Submission #236445

#	Time	Username	Problem	Language	Result	Execution time	Memory
236445		Charis02	Mergers (JOI19_mergers)	C++14	0 / 100	219 ms	142444 KiB

mergers

#include<iostream>
#include<stdio.h>
#include<vector>
#include<cmath>
#include<queue>
#include<string.h>
#include<map>
#include<set>
#include<algorithm>
#define ll long long
#define pi pair < ll,ll >
#define mp(a,b) make_pair(a,b)
#define mid (low+high)/2
#define rep(i,a,b) for(int i = a;i < b;i++)
#define N 500004
#define INF 1e9+7

using namespace std;

struct dsu{
    ll par[N];
    ll sz[N];
    ll col[N];

    ll get(ll x)
    {
        return (par[x] == x) ? x : get(par[x]);
    }

    void join(ll x,ll y)
    {
        ll fx=get(x);
        ll fy = get(y);
        if(fx==fy)
            return;
        if(sz[fx]<sz[fy])
            swap(fx,fy);
        sz[fx]+=sz[fy];
        par[fy]=fx;
    }

    ll get_col(ll x)
    {
        return col[get(x)];
    }
};

struct segment_tree{
    ll val[4*N];

    void upd(ll low,ll high,ll pos,ll slow,ll v)
    {
        if(low==high&&low==slow)
        {
            val[pos]+=v;
            return;
        }
        if(low>slow||high<slow)
            return;
        upd(low,mid,pos*2+1,slow,v);
        upd(mid+1,high,pos*2+2,slow,v);
        val[pos]=val[pos*2+1]+val[pos*2+2];
        return;
    }

    ll query(ll low,ll high,ll pos,ll slow,ll shigh)
    {
        if(low>=slow&&high<=shigh)
            return val[pos];
        if(low>shigh||high<slow)
            return 0;
        return query(low,mid,pos*2+1,slow,shigh)+query(mid+1,high,pos*2+2,slow,shigh);
    }
};

ll n,k,answer;
bool processed[N];
vector < ll > graph[N];
vector < ll > cities[N];
vector < ll > queries[N];
ll cnt = 0;
ll colors;
ll ar[N],deg[N];
ll dp[N][20],tin[N],tout[N];
segment_tree seg;
dsu comp;

ll get_par(ll x,ll i)
{
    if(dp[x][i] !=-1)
        return dp[x][i];
    return dp[x][i] = get_par(get_par(x,i-1),i-1);
}

bool is_ancestor(ll x,ll y)
{
    return (tin[x] <= tin[y] && tout[x] >= tout[y]);
}

ll get_lca(ll x,ll y)
{
    if(is_ancestor(x,y))
        return x;
    if(is_ancestor(y,x))
        return y;

    for(int i = 19;i >= 0;i--)
    {
        if(!is_ancestor(get_par(x,i),y))
            x = get_par(x,i);
    }

    return get_par(x,0);
}

void init(ll cur,ll par)
{
    dp[cur][0] = par;
    tin[cur] = cnt++;

    rep(i,0,graph[cur].size())
    {
        ll v = graph[cur][i];
        if(v==par)
            continue;
        init(v,cur);
    }

    tout[cur]=cnt;
    return;
}

void dfs(ll cur,ll par)
{
    rep(i,0,queries[cur].size())
        seg.upd(0,n,0,tin[queries[cur][i]],1);

    rep(i,0,graph[cur].size())
    {
        ll v= graph[cur][i];

        if(v==par)
            continue;

        if(seg.query(0,n,0,tin[v],tout[v])!=0)
            comp.join(cur,v);
        dfs(v,cur);
    }

    rep(i,0,queries[cur].size())
        seg.upd(0,n,0,tin[queries[cur][i]],-1);
}

void process_paths()
{
    rep(s,1,k+1)
    {
        ll x = cities[s][0];

        rep(j,1,cities[s].size())
        {
            ll y = cities[s][j];
            ll lca = get_lca(x,y);

            if(lca==x)
                queries[x].push_back(y);
            else if(lca==y)
                queries[y].push_back(x);
            else
            {
                queries[lca].push_back(x);
                queries[lca].push_back(y);
            }
        }
    }

    dfs(1,1);

    rep(i,1,n+1)
    {
        if(comp.get_col(i) == 0)
            comp.col[comp.get(i)] = ++colors;
        ar[i] = comp.get_col(i);
    }

    return;
}

void calc()
{
    rep(i,1,n+1)
    {
        rep(j,0,graph[i].size())
        {
            ll v = graph[i][j];
            if(ar[i] != ar[v])
                deg[ar[i]]++;
        }
    }

    ll leaves = 0;

    rep(i,1,colors)
    {
        if(deg[i] == 1)
            leaves++;
    }

    answer = (leaves+1)/2;

    return;
}

int main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(0);

    cin >> n >> k;

    rep(i,0,n-1)
    {
        ll a,b;
        cin >> a >> b;
        graph[a].push_back(b);
        graph[b].push_back(a);
    }

    rep(i,1,n+1)
    {
        ll s;
        comp.par[i]=i;
        comp.sz[i]=1;
        cin >> s;
        cities[s].push_back(i);
    }

    memset(seg.val,0,sizeof seg.val);
    memset(dp,-1,sizeof dp);
    init(1,1);
    process_paths();
    calc();

    cout << answer << endl;

    return 0;
}

/*
5 4
1 2
2 3
3 4
4 5
1
2
3
4
1

*/

Compilation message (stderr)

mergers.cpp: In function 'void init(long long int, long long int)':
mergers.cpp:14:36: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
 #define rep(i,a,b) for(int i = a;i < b;i++)
mergers.cpp:121:9:
     rep(i,0,graph[cur].size())
         ~~~~~~~~~~~~~~~~~~~~~       
mergers.cpp:121:5: note: in expansion of macro 'rep'
     rep(i,0,graph[cur].size())
     ^~~
mergers.cpp: In function 'void dfs(long long int, long long int)':
mergers.cpp:14:36: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
 #define rep(i,a,b) for(int i = a;i < b;i++)
mergers.cpp:135:9:
     rep(i,0,queries[cur].size())
         ~~~~~~~~~~~~~~~~~~~~~~~     
mergers.cpp:135:5: note: in expansion of macro 'rep'
     rep(i,0,queries[cur].size())
     ^~~
mergers.cpp:14:36: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
 #define rep(i,a,b) for(int i = a;i < b;i++)
mergers.cpp:138:9:
     rep(i,0,graph[cur].size())
         ~~~~~~~~~~~~~~~~~~~~~       
mergers.cpp:138:5: note: in expansion of macro 'rep'
     rep(i,0,graph[cur].size())
     ^~~
mergers.cpp:14:36: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
 #define rep(i,a,b) for(int i = a;i < b;i++)
mergers.cpp:150:9:
     rep(i,0,queries[cur].size())
         ~~~~~~~~~~~~~~~~~~~~~~~     
mergers.cpp:150:5: note: in expansion of macro 'rep'
     rep(i,0,queries[cur].size())
     ^~~
mergers.cpp: In function 'void process_paths()':
mergers.cpp:14:36: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
 #define rep(i,a,b) for(int i = a;i < b;i++)
mergers.cpp:160:13:
         rep(j,1,cities[s].size())
             ~~~~~~~~~~~~~~~~~~~~    
mergers.cpp:160:9: note: in expansion of macro 'rep'
         rep(j,1,cities[s].size())
         ^~~
mergers.cpp: In function 'void calc()':
mergers.cpp:14:36: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
 #define rep(i,a,b) for(int i = a;i < b;i++)
mergers.cpp:193:13:
         rep(j,0,graph[i].size())
             ~~~~~~~~~~~~~~~~~~~     
mergers.cpp:193:9: note: in expansion of macro 'rep'
         rep(j,0,graph[i].size())
         ^~~

• Subtask 1: 0 / 10
• Subtask 2: 0 / 24
• Subtask 3: 0 / 14
• Subtask 4: 0 / 22
• Subtask 5: 0 / 30