Submission #889531

#	Time	Username	Problem	Language	Result	Execution time	Memory
889531		zeta7532	Railway (BOI17_railway)	C++17	0 / 100	449 ms	53564 KiB

railway

#include <bits/stdc++.h>
#pragma GCC target("avx2")
#pragma GCC optimize("O3")
#pragma GCC optimize("unroll-loops")
using namespace std;
using ll = long long;
using Graph = vector<vector<ll>>;
using P = pair<ll,ll>;
const ll mod = 998244353;
#define fi first
#define se second
#define rep(i,n) for(ll i=0;i<n;i++)
#define all(x) x.begin(),x.end()
#define faster ios::sync_with_stdio(false);cin.tie(nullptr)

template<typename T>
struct RMQ1{
    const T INF=1LL<<60;
    ll N;
    vector<P> dat;
    RMQ1(ll N_):N(),dat(N_*4){
        ll x=1;
        while(N_>x) x*=2;
        N=x;
    }
    void update(ll i,T x,T y){
        i+=N-1;
        dat[i].fi=x;
        dat[i].se=y;
        while(i>0){
            i=(i-1)/2;
            if(dat[i*2+1].se<=dat[i*2+2].se){
                dat[i].fi=dat[i*2+1].fi;
                dat[i].se=dat[i*2+1].se;
            }
            if(dat[i*2+1].se>dat[i*2+2].se){
                dat[i].fi=dat[i*2+2].fi;
                dat[i].se=dat[i*2+2].se;
            }
        }
    }
    P query(ll a,ll b){return query_sub(a,b,0,0,N);}
    P query_sub(ll a,ll b,ll k,ll l,ll r){
        if(r<=a||b<=l){
            return {INF,INF};
        }else if(a<=l&&r<=b){
            return dat[k];
        }else{
            P vl=query_sub(a,b,k*2+1,l,(l+r)/2);
            P vr=query_sub(a,b,k*2+2,(l+r)/2,r);
            P q;
            if(vl.se<=vr.se){
                q.fi=vl.fi;
                q.se=vl.se;
            }
            if(vl.se>vr.se){
                q.fi=vr.fi;
                q.se=vr.se;
            }
            return q;
        }
    }
};

template<typename T>
struct RMQ2{
    const T INF=0;
    ll N;
    vector<T> dat;
    RMQ2(ll N_):N(),dat(N_*4){
        ll x=1;
        while(N_>x) x*=2;
        N=x;
    }
    void update(ll i,T x){
        i+=N-1;
        dat[i]+=x;
        while(i>0){
            i=(i-1)/2;
            dat[i]=dat[i*2+1]+dat[i*2+2];
        }
    }
    T query(ll a,ll b){return query_sub(a,b,0,0,N);}
    T query_sub(ll a,ll b,ll k,ll l,ll r){
        if(r<=a||b<=l){
            return INF;
        }else if(a<=l&&r<=b){
            return dat[k];
        }else{
            T vl=query_sub(a,b,k*2+1,l,(l+r)/2);
            T vr=query_sub(a,b,k*2+2,(l+r)/2,r);
            return vl+vr;
        }
    }
};

ll cl[100500],cr[100500];
ll Order=0;
ll dis[100500];
ll par[100500];

void dfs(const Graph &G,ll pos,ll dep,ll pre){
    dis[pos]=dep;
    Order++;
    cl[pos]=Order;
    for(ll to:G[pos]){
        if(to==pre) continue;
        par[to]=pos;
        dfs(G,to,dep+1,pos);
    }
    Order++;
    cr[pos]=Order;
}

struct centroid_decomposition{
    ll N;
    vector<vector<ll>> G;
    vector<ll> sz;
    vector<ll> par;
    vector<ll> dep;
    centroid_decomposition(const vector<vector<ll>> &G):N(G.size()),G(G),sz(N),par(N,-1),dep(N,-1) {build(0,-1,0);}
    void calcSize(ll u,ll p){
        sz[u]=1;
        for(ll v:G[u]){
            if(dep[v]==-1&&v!=p){
                calcSize(v,u);
                sz[u]+=sz[v];
            }
        }
    }
    void build(ll u,ll p,ll d){
        calcSize(u,-1);
        ll sum=sz[u];
        bool ok=false;
        ll pp=-1;
        while(!ok){
            ok=true;
            for(ll v:G[u]){
                if(dep[v]==-1&&v!=pp&&2*sz[v]>sum){
                    pp=u;
                    u=v;
                    ok=false;
                    break;
                }
            }
        }
        par[u]=p;
        dep[u]=d;
        for(ll v:G[u]){
            if(dep[v]==-1){
                build(v,u,d+1);
            }
        }
    }
    pair<vector<vector<ll>>,pair<vector<ll>,ll>> cdtree(){
        vector<vector<ll>> res(N);
        ll r;
        rep(i,N){
            if(par[i]!=-1) res[par[i]].push_back(i);
            else r=i;
        }
        return {res,{par,r}};
    }
};

vector<ll> bfs(vector<vector<ll>> G,ll s){
    ll N=G.size();
    vector<ll> dis(N,-1);
    dis[s]=0;
    queue<ll> que;
    que.push(s);
    while(!que.empty()){
        ll v=que.front();
        que.pop();
        for(ll nv:G[v]){
            if(dis[nv]!=-1) continue;
            dis[nv]=dis[v]+1;
            que.push(nv);
        }
    }
    return dis;
}

struct LCA{
    vector<vector<ll>> par;
    vector<ll> dis;
    LCA(const Graph &G,ll root=0){init (G,root);}
    void init(const Graph &G,ll root=0){
        ll v=G.size();
        ll k=1;
        while((1<<k)<v) k++;
        par.assign(k,vector<ll>(v,-1));
        dis.assign(v,-1);
        dfs(G,root,-1,0);
        for(ll i=0;i+1<k;i++){
            for(ll j=0;j<v;j++){
                if(par[i][j]<0) par[i+1][j]=-1;
                else par[i+1][j]=par[i][par[i][j]];
            }
        }
    }
    void dfs(const Graph &G,ll v,ll p,ll d){
        par[0][v]=p;
        dis[v]=d;
        for(ll nv:G[v]){
            if(nv!=p) dfs(G,nv,v,d+1);
        }
    }
    ll query(ll u,ll v){
        if(dis[u]<dis[v]) swap(u,v);
        ll K=par.size();
        for(ll k=0;k<K;k++){
            if((dis[u]-dis[v])>>k&1){
                u=par[k][u];
            }
        }
        if(u==v) return u;
        for(ll k=K-1;k>=0;k--){
            if(par[k][u]!=par[k][v]){
                u=par[k][u];
                v=par[k][v];
            }
        }
        return par[0][u];
    }
    ll query_prev(ll v,ll p){
        ll K=par.size();
        for(ll k=0;k<K;k++){
            if((dis[v]-dis[p]-1)>>k&1){
                v=par[k][v];
            }
        }
        return v;
    }
};

int main() {
    faster;
    ll N,Q,K;
    cin >> N >> Q >> K;
    vector<ll> A(N-1),B(N-1);
    vector<vector<ll>> G(N);
    rep(i,N-1) cin >> A[i] >> B[i];
    rep(i,N-1) A[i]--,B[i]--;
    rep(i,N-1) G[A[i]].push_back(B[i]),G[B[i]].push_back(A[i]);
    dfs(G,0,0,-1);
    vector<ll> inv(N,-1);
    rep(i,N-1){
        if(dis[A[i]]>dis[B[i]]) inv[A[i]]=i;
        if(dis[B[i]]>dis[A[i]]) inv[B[i]]=i;
    }
    RMQ2<ll> SUM(200500);
    centroid_decomposition CD(G);
    vector<vector<ll>> res=CD.cdtree().fi;
    vector<ll> PAR=CD.cdtree().se.fi;
    ll r=CD.cdtree().se.se;
    vector<ll> DIS=bfs(res,r);
    LCA ans(G);
    while(Q--){
        ll S;
        cin >> S;
        set<pair<ll,ll>> s;
        rep(i,S){
            ll t;
            cin >> t;
            t--;
            s.insert({-DIS[t],t});
        }
        while(1){
            auto it=*s.begin();
            if(it.fi==0) break;
            s.erase(it);
            ll x=it.se;
            ll y=PAR[x];
            s.insert({-DIS[y],y});
            ll lca=ans.query(x,y);
            if(x!=lca){
                ll x_goal=ans.query_prev(x,lca);
                SUM.update(cl[x],1);
                SUM.update(cr[x_goal],-1);
            }
            if(y!=lca){
                ll y_goal=ans.query_prev(y,lca);
                SUM.update(cl[y],1);
                SUM.update(cr[y_goal],-1);
            }
        }
    }
    vector<ll> ANS;
    for(ll i=1;i<N;i++){
        ll cnt=SUM.query(1,cl[i]+1)-SUM.query(1,cl[PAR[i]]+1);
        if(cnt>=K) ANS.push_back(inv[i]+1);
    }
    sort(all(ANS));
    cout << ANS.size() << endl;
    rep(i,ANS.size()) cout << ANS[i] << " ";
    cout << endl;
    return 0;
}

Compilation message (stderr)

railway.cpp: In function 'int main()':
railway.cpp:12:30: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   12 | #define rep(i,n) for(ll i=0;i<n;i++)
......
  296 |     rep(i,ANS.size()) cout << ANS[i] << " ";
      |         ~~~~~~~~~~~~          
railway.cpp:296:5: note: in expansion of macro 'rep'
  296 |     rep(i,ANS.size()) cout << ANS[i] << " ";
      |     ^~~

• Subtask 1: 0 / 8
• Subtask 2: 0 / 15
• Subtask 3: 0 / 7
• Subtask 4: 0 / 29
• Subtask 5: 0 / 16
• Subtask 6: 0 / 25