Submission #1096573

#	Time	Username	Problem	Language	Result	Execution time	Memory
1096573		AlgorithmWarrior	Mousetrap (CEOI17_mousetrap)	C++14	100 / 100	879 ms	158548 KiB

mousetrap

#include <bits/stdc++.h>
#define MAX 1000005

using namespace std;

vector<int>tree[MAX];
bool on_path[MAX];
int dp[MAX];
int h[MAX];
int tata[MAX];

int n,t,m;

void read(){
    cin>>n>>t>>m;
    int i;
    for(i=1;i<n;++i){
        int a,b;
        cin>>a>>b;
        tree[a].push_back(b);
        tree[b].push_back(a);
    }
}

void dfs(int nod){
    for(auto fiu : tree[nod])
    if(fiu!=tata[nod]){
        h[fiu]=h[nod]+1;
        tata[fiu]=nod;
        dfs(fiu);
    }
}

void dfs_dp(int nod,int tata){
    int first_val=0,second_val=0;
    for(auto fiu : tree[nod])
    if(fiu!=tata && !on_path[fiu]){
        dfs_dp(fiu,nod);
        ++dp[nod];
        if(dp[fiu]>first_val){
            second_val=first_val;
            first_val=dp[fiu];
        }
        else
            if(dp[fiu]>second_val)
                second_val=dp[fiu];
    }
    dp[nod]+=second_val;
}

int lca(int nod1,int nod2){
    while(nod1!=nod2){
        if(h[nod1]>h[nod2])
            nod1=tata[nod1];
        else
            nod2=tata[nod2];
    }
    return nod1;
}

vector<int>path(int nod1,int nod2){
    int lc=lca(nod1,nod2);
    vector<int>pth1,pth2;
    while(nod1!=lc){
        nod1=tata[nod1];
        pth1.push_back(nod1);
    }
    while(nod2!=lc){
        pth2.push_back(nod2);
        nod2=tata[nod2];
    }
    reverse(pth2.begin(),pth2.end());
    for(auto el : pth2){
        pth1.push_back(el);
    }
    return pth1;
}

int v[MAX];
int add[MAX];
int order[MAX];
int order_rev[MAX];
int sum[MAX];
int psz;
int ind;

bool check_sum(){
    int i;
    int s=0;
    for(i=1;i<=psz;++i){
        s+=sum[i];
        if(s>i)
            return 0;
    }
    return 1;
}

void get_son(int nod){
    for(auto fiu : tree[nod]){
        if(!on_path[fiu]){
            v[++ind]=dp[fiu];
        }
    }
}

void init_dfs_dp(){
    int nod1=m,nod2=t;
    on_path[m]=on_path[t]=1;
    while(nod1!=nod2){
        if(h[nod1]>h[nod2]){
            nod1=tata[nod1];
            on_path[nod1]=1;
        }
        else{
            nod2=tata[nod2];
            on_path[nod2]=1;
        }
    }
    int i,j;
    for(i=1;i<=n;++i)
        if(on_path[i])
            dfs_dp(i,0);
    vector<int>pth=path(t,m);
    for(i=0;i<pth.size();++i){
        get_son(pth[i]);
        add[i+1]=ind;
        for(j=add[i]+1;j<=add[i+1];++j){
            v[j]+=add[i+1];
            order[j]=i+1;
        }
    }
    for(i=1;i<=ind;++i)
        order_rev[i]=pth.size()+1-order[i];
    psz=pth.size();
}

int get_max(){
    int i;
    int id=0;
    for(i=1;i<=ind;++i)
        if(v[id]<v[i])
            id=i;
    return id;
}

int solve(){
    int rasp=0;
    int luate=0;
    while(check_sum() && luate<ind){
        ++luate;
        int id=get_max();
        rasp=v[id];
        v[id]=-2e9;
        int idcur=order[id];
        int i;
        for(i=1;i<=add[idcur-1];++i)
            ++v[i];
        int idcurrev=order_rev[id];
        ++sum[idcurrev];
    }
    if(check_sum())
        return ind;
    return rasp;
}

void write(){
    cout<<solve();
}

int main()
{
    read();
    dfs(1);
    init_dfs_dp();
    write();
    return 0;
}

Compilation message (stderr)

mousetrap.cpp: In function 'void init_dfs_dp()':
mousetrap.cpp:124:14: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  124 |     for(i=0;i<pth.size();++i){
      |             ~^~~~~~~~~~~

• Subtask 1: 20 / 20
• Subtask 2: 25 / 25
• Subtask 3: 20 / 20
• Subtask 4: 35 / 35