Submission #1281128

#	Time	Username	Problem	Language	Result	Execution time	Memory
1281128		longdeptrai	Railway (BOI17_railway)	C++20	100 / 100	133 ms	38840 KiB

railway

#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;
#define LongDepTrai "railway"
#define ll long long
#define int long long
#define ull unsigned long long
#define ld long double
#define ii pair<int,int>
#define iii pair<int,ii>
#define iv pair<ii,ii>
#define pll pair<ll,ll>
#define vi vector<int>
#define vii vector<ii>
#define vll vector<ll>
#define fi first
#define se second
#define pb push_back
#define all(x) (x).begin(), (x).end()
#define sz(x) int((x).size())
#define order_set(T) tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>

inline ll add(ll a, ll b, ll mod){ a += b; if(a >= mod) a -= mod; return a; }
inline ll sub(ll a, ll b, ll mod){ a -= b; if(a < 0) a += mod; return a; }
inline ll mul(ll a, ll b, ll mod){ return ( (ll)a * b ) % mod; }

static mt19937_64 rng((unsigned)chrono::steady_clock::now().time_since_epoch().count());

const int N=2e5+9;
const int lg=18;
int n, test, a[N], k,par[N][lg+2], tour[N], timer=0,timer2=0, h[N], kval, dist[N], sz[N], tin[N], tout[N],heavy[N],head[N],pos[N];
//vector<ii> v[N];
vector<int> g[N];
int cnt = 0;
bool ok[N];
struct BIT {
    int n;
    vector<long long> bit; // 1-indexed

    BIT(int n=0) { init(n); }

    void init(int _n) {
        n = _n;
        bit.assign(n+1, 0);
    }

    // add `val` to index idx (point update on internal BIT)
    void addPoint(int idx, long long val) {
        for (; idx <= n; idx += idx & -idx) bit[idx] += val;
    }

    // prefix sum up to idx (inclusive)
    long long prefixSum(int idx) const {
        long long s = 0;
        for (; idx > 0; idx -= idx & -idx) s += bit[idx];
        return s;
    }

    // RANGE UPDATE: add val to all positions in [l, r]
    // NOTE: l and r must be 1-indexed and 1 <= l <= r <= n
    void rangeAdd(int l, int r, long long val) {
        if (l > r) return;
        addPoint(l, val);
        if (r+1 <= n) addPoint(r+1, -val);
    }

    // POINT QUERY: get value at position pos (1-indexed)
    long long pointQuery(int pos) const {
        return prefixSum(pos);
    }
} bit;

void dfs(int u, int p) {
	tin[u] = ++timer;
	for (int e : g[u]) {
		int x = e;
		if (x == p) continue;
		h[x] = h[u] + 1;
		//dist[x] = dist[u] + y;
		par[x][0] = u;
		dfs(x, u);
	}
	tout[u] = timer;
}
int dfs2(int u,int p){
    int mx_sz=0;
    int cur_sz=1;

    for(int x:g[u]){
        if(x==p) continue;
      //  h[x]=h[u]+1;
     //   par[x]=u;
        int son_sz=dfs2(x,u);
        if(son_sz>mx_sz){
            mx_sz=son_sz;
            heavy[u]=x;
        }
        cur_sz+=son_sz;
    }
    return cur_sz;
}
void hld(int u,int h){
    head[u]=h;
    pos[u]=++timer2;
    tour[timer2]=u;
    if(heavy[u]) hld(heavy[u],h);
    for(int x:g[u]){
        if(x!=par[u][0] && x!=heavy[u]){
            hld(x,x);
        }
    }
}
void change(int x,int y){
    while(head[x]!=head[y]){
        if(h[head[x]]<h[head[y]]) swap(x,y);
       // Update(1,pos[head[x]],pos[x],1,n,k);
        bit.rangeAdd(pos[head[x]],pos[x],1);
    //    cout<<tour[pos[head[x]]]<<" "<<x<<"\n";
        x=par[head[x]][0];
    }
    if(h[y]<h[x]) swap(x,y);
 //   Update(1,pos[x]+1,pos[y],1,n,k);
    bit.rangeAdd(pos[x]+1,pos[y],1);
  //  cout<<tour[pos[x]+1]<<" "<<tour[pos[y]]<<"\n";
}

int lca(int u, int v) {
	if (h[u] < h[v]) swap(u, v);
	for (int i = lg; i >= 0; i--) {
		if ((h[u] - h[v]) >= (1ll << i)) u = par[u][i];
	}
	if (u == v) return u;
	for (int i = lg; i >= 0; i--) {
		if (par[u][i] != par[v][i]) {
			u = par[u][i];
			v = par[v][i];
		}
	}
	return par[u][0];
}
void prelca() {
	dfs(1, -1);
	for (int j = 1; j <= lg; j++) {
		for (int i = 1; i <= n; i++) {
			par[i][j] = par[par[i][j-1]][j-1];
		}
	}
}
bool ss(int x, int y) {
	return tin[x] < tin[y];
}
bool inside(int x, int y) {
	return (tin[x] <= tin[y] && tin[y] <= tout[x]);
}
int build(vector<int> &ver) {
	sort(ver.begin(), ver.end(), ss);
//	for(int x:ver) cout<<x<<" ";
//	cout<<"\n";
	int p = ver.size();
	for (int i = 0; i < p-1; i++) {
		int z = lca(ver[i], ver[i+1]);
		ver.pb(z);
	}
	sort(ver.begin(), ver.end(), ss);
	ver.erase(unique(ver.begin(), ver.end()), ver.end());

	vector<int> st;
	st.pb(ver[0]);
	for (int i = 1; i < ver.size(); i++) {
		int u = ver[i];
		while(!st.empty() && !inside(st.back(), u)) st.pop_back();
		assert(!st.empty());
		int last = st.back();
		//cout << 0<<" "<<last << " " << u << endl;
		change(last,u);

		//g[last].pb(u);
		st.pb(u);
	}
	return ver[0];
}
void solve(){
    cin >> n >> test>>kval;
    bit.init(n+2);
    vii ed;
    for(int i=1;i<n;i++){
        int u,v;
        cin>>u>>v;
        g[u].pb(v);
        g[v].pb(u);
        ed.pb({u,v});
        //v[u].pb(v);

    }
	prelca();
	dfs2(1,-1);
    hld(1,1);
	for (int Q = 1; Q <= test; Q++) {
		cin >> k;
		vector<int> ver(k);
		for (int i = 0; i < k; i++) {
			cin >> ver[i];
            cnt += ver[i];
//            cout << "! " << ver[i] << " ";
			sz[ver[i]] += ver[i];
		}
		int goc = build(ver);
//        for(int i=2;i<=n;i++){
//            cout<<bit.pointQuery(pos[i])<<" ";
//        }
//        cout<<"\n";
	}
	vi res;
	for(int i=0;i<ed.size();i++){
        int u=ed[i].fi;
        int v=ed[i].se;
        if(h[u]<h[v]) swap(u,v);
        if(bit.pointQuery(pos[u])>=kval){
            res.pb(i+1);
        }
	}
    cout<<res.size()<<"\n";
    for(int x:res) cout<<x<<" ";
}
signed main(){
    ios_base::sync_with_stdio(0);
    cin.tie(0);cout.tie(0);

    int t=1;
    while(t--){
        solve();
    }
}

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