Submission #824972

#	Time	Username	Problem	Language	Result	Execution time	Memory
824972		Moses	Road Closures (APIO21_roads)	C++14	100 / 100	371 ms	30100 KiB

roads

#include "roads.h"
#include <bits/stdc++.h>
#define F first
#define S second
#define pb push_back
//#define int long long
using namespace std;
const int N = 1e5+7;
const int MAXW = 1e9+7;
typedef pair<int,int> pii;
typedef long long ll;
vector<pii> g[N];
int deg[N];

mt19937 RNG(chrono::steady_clock::now().time_since_epoch().count());
template<typename A, typename B> ostream& operator<<(ostream &os, const pair<A, B> &p) { return os << '(' << p.first << ", " << p.second << ')'; }
template<typename T_container, typename T = typename enable_if<!is_same<T_container, string>::value, typename T_container::value_type>::type> ostream& operator<<(ostream &os, const T_container &v) { os << '{'; string sep; for (const T &x : v) os << sep << x, sep = ", "; return os << '}'; }
void dbg_out() { cerr << endl; }
template<typename Head, typename... Tail> void dbg_out(Head H, Tail... T) { cerr << ' ' << H; dbg_out(T...); }
#ifdef DEBUG
#define dbg(...) cerr << "(" << #__VA_ARGS__ << "):", dbg_out(__VA_ARGS__)
#else
#define dbg(...)
#endif
struct Treap {
	int data, priority;
	int mxdata;
	array<Treap*, 2> kids;
	int subtreeSize;
	ll sum;

	Treap(int _data);
};

int size(Treap *me) {
	return me == NULL ? 0 : me->subtreeSize;
}

int mxdata(Treap *me) {
	return me == NULL ? 0 : me->mxdata;
}

ll sum(Treap *me) {
	return me == NULL ? 0 : me->sum;
}

void recalc(Treap *me) {
	if (me==NULL) return;
	me->subtreeSize = 1;
	me->sum = me->data;
	me->mxdata = me->data;
	for (Treap* t:me->kids) if (t != NULL) me->subtreeSize += t->subtreeSize;
	for (Treap* t:me->kids) if (t != NULL) me->sum += t->sum;
	for (Treap* t:me->kids) if (t != NULL) me->mxdata = (me->mxdata > t->mxdata ? me->mxdata : t->mxdata);
}

Treap* merge(Treap *l, Treap *r) {
	if (l==NULL) return r;
	if (r==NULL) return l;
	if (l->priority < r->priority) {
		l->kids[1]=merge(l->kids[1], r);
		recalc(l);
		return l;
	}
	else {
		r->kids[0]=merge(l, r->kids[0]);
		recalc(r);
		return r;
	}
}

array<Treap*, 2> split(Treap *me, int nInLeft) {
	if (me == NULL) return {NULL, NULL};
	if (size(me->kids[0])>=nInLeft) {
		array<Treap*, 2> leftRes=split(me->kids[0], nInLeft);
		me->kids[0]=leftRes[1];
		recalc(me);
		return {leftRes[0], me};
	}
	else {
		nInLeft = nInLeft - size(me->kids[0]) - 1;
		array<Treap*, 2> rightRes = split(me->kids[1], nInLeft);
		me->kids[1] = rightRes[0];
		recalc(me);
		return {me, rightRes[1]};
	}
	return {NULL, NULL};
}

array<Treap*, 2> split_on_value(Treap *me, ll nInLeft) {
	if (me == NULL) return {NULL, NULL};
	//prop(me);
	if (mxdata(me->kids[0])>nInLeft) {
		array<Treap*, 2> leftRes=split_on_value(me->kids[0], nInLeft);
		me->kids[0]=leftRes[1];
		recalc(me);
		return {leftRes[0], me};
	}
	else if (me->data <= nInLeft) {
		array<Treap*, 2> rightRes = split_on_value(me->kids[1], nInLeft);
		me->kids[1] = rightRes[0];
		recalc(me);
		return {me, rightRes[1]};
	} else {
		array<Treap*, 2> leftRes=split_on_value(me->kids[0], nInLeft);
		me->kids[0]=leftRes[1];
		recalc(me);
		return {leftRes[0], me};
	}
	return {NULL, NULL};
}

Treap::Treap(int _data) {
	kids={NULL, NULL};
	this->data = _data;
	recalc(this);
	this->priority = (int)RNG();
}

pair<int,ll> query2(Treap *&me, int cnt) {
	array<Treap*, 2> tmp = split(me,cnt);
	pair<int,ll> ret = {size(tmp[0]),sum(tmp[0])};
	me = merge(tmp[0],tmp[1]);
	return ret;
}

pair<int,ll> query(Treap *&me, ll L,ll R,int cnt) {
	array<Treap*, 2> tmp = split_on_value(me,L-1);
	array<Treap*, 2> tmp2 = split_on_value(tmp[1],R);
	
	pair<int,ll> ret = query2(tmp2[0], cnt);
	
	tmp[1] = merge(tmp2[0],tmp2[1]);
	me = merge(tmp[0],tmp[1]);
	return ret;
}

void insert(Treap *&me,int x) {
	array<Treap*, 2> tmp = split_on_value(me,x);
	Treap* tmp2 = new Treap(x);
	tmp[0] = merge(tmp[0],tmp2);
	me = merge(tmp[0],tmp[1]);
}

Treap* tr[N];
int k;
bool vis[N];
ll dp[N][2];
void dfs(int u,int p = -1) {
	vis[u] = 1;
	vector<pii> kids;
	dp[u][0] = 0;
	for(auto [v,w]:g[u]) {
		if (v == p) continue;
		dfs(v,u);
		kids.pb({v,w});
		dp[u][0] += dp[v][0];
	}
	dp[u][1] = dp[u][0];
	sort(kids.begin(),kids.end(),[&](pii u,pii v) {return dp[u.F][1]-dp[u.F][0]+u.S < dp[v.F][1]-dp[v.F][0]+v.S;});
	
	int need = max(0,deg[u]-k);
	//dbg(u,need);
	int last = -1;
	for(auto [v,w]:kids) {
		if (dp[v][1]-dp[v][0]+w < 0) {
			need = max(need-1,0);
			last = dp[v][1]-dp[v][0]+w;
			dp[u][0] += dp[v][1]-dp[v][0]+w;
			continue;
		}
		if (need == 0) continue;
		auto ret = query(tr[u],last+1,dp[v][1]-dp[v][0]+w,need);
		last = dp[v][1]-dp[v][0]+w;
		need -= ret.F;
		assert(need >= 0);
		dp[u][0] += ret.S;
		//dbg(u,need,dp[u][0],last+1,w+dp[v][1]-dp[v][0]);
		if (need == 0) continue;
		need--;
		assert(need >= 0);
		dp[u][0] += dp[v][1]-dp[v][0]+w;
	}
	if (need > 0) {
		auto ret = query(tr[u],last+1,MAXW,need);
		dbg(k,u,ret);
		dp[u][0] += ret.S;
	}
	int need1 = max(0,deg[u]-k-1);
	last = -1;
	for(auto [v,w]:kids) {
		if (dp[v][1]-dp[v][0]+w < 0) {
			need1 = max(need1-1,0);
			last = dp[v][1]-dp[v][0]+w;
			dp[u][1] += dp[v][1]-dp[v][0]+w;
			continue;
		}
		if (need1 == 0) continue;
		auto ret = query(tr[u],last+1,dp[v][1]-dp[v][0]+w,need1);
		last = dp[v][1]-dp[v][0]+w;
		need1 -= ret.F;
		assert(need1 >= 0);
		dp[u][1] += ret.S;
		if (need1 == 0) continue;
		need1--;
		assert(need1 >= 0);
		dp[u][1] += dp[v][1]-dp[v][0]+w;
	}
	if (need1 > 0) {
		auto ret = query(tr[u],last+1,MAXW,need1);
		dp[u][1] += ret.S;
	}
	dbg(k,u,dp[u][0],dp[u][1]);
}

vector<long long> minimum_closure_costs(int N, vector<int> U,
                                             vector<int> V,
                                             vector<int> W) {
	
	for(int i = 0 ; i < N-1 ; i++) {
		g[U[i]].pb({V[i],W[i]});
		g[V[i]].pb({U[i],W[i]});
		deg[U[i]]++;
		deg[V[i]]++;
	}
	
	vector<int> vec;
	
	for(int i = 0 ; i < N ; i++) {
		vec.pb(i);
		sort(g[i].begin(),g[i].end(),[&](pii u,pii v) {return deg[u.F] > deg[v.F];});
	}
	
	sort(vec.begin(),vec.end(),
	[](int u,int v) {return deg[u] < deg[v];});
	
	int ptr = 0;
	vector<ll> ans(N,0);
	for(k = 0 ; k < N; k++) {
		while(ptr < N && deg[vec[ptr]] <= k) ptr++;
		dbg(k,ptr);
		for(int i = ptr ; i < N ; i++) {
			int u = vec[i];
			//dbg(k,u);
			while(!g[u].empty() && deg[g[u].back().F] <= k) {
				int w = g[u].back().S;
				g[u].pop_back();
				//dbg(query2(tr[u],1));
				insert(tr[u],w);
				//dbg(query2(tr[u],1));
				dbg(k,u,w);
			}
		}
		
		for(int i = ptr ; i < N ; i++) {
			int u = vec[i];
			if (!vis[u]) {dfs(u);ans[k] += dp[u][0];}
		}
		
		for(int i = ptr ; i < N ; i++) {
			int u = vec[i];
			vis[u] = 0;
		}
	}
	return ans;
}

Compilation message (stderr)

roads.cpp: In function 'void dfs(int, int)':
roads.cpp:153:11: warning: structured bindings only available with '-std=c++17' or '-std=gnu++17'
  153 |  for(auto [v,w]:g[u]) {
      |           ^
roads.cpp:165:11: warning: structured bindings only available with '-std=c++17' or '-std=gnu++17'
  165 |  for(auto [v,w]:kids) {
      |           ^
roads.cpp:191:11: warning: structured bindings only available with '-std=c++17' or '-std=gnu++17'
  191 |  for(auto [v,w]:kids) {
      |           ^

• Subtask 1: 5 / 5
• Subtask 2: 7 / 7
• Subtask 3: 14 / 14
• Subtask 4: 10 / 10
• Subtask 5: 17 / 17
• Subtask 6: 25 / 25
• Subtask 7: 22 / 22