Submission #444780

#	Time	Username	Problem	Language	Result	Execution time	Memory
444780		arnold518	Road Closures (APIO21_roads)	C++17	100 / 100	802 ms	359612 KiB

roads

#include "roads.h"
#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

const int MAXN = 1e5;

int N;
vector<pii> adj[MAXN+10], adj2[MAXN+10];
int par[MAXN+10], val[MAXN+10], sz[MAXN+10], deg[MAXN+10], dfn[MAXN+10];
ll ans[MAXN+10];

struct Edge
{
	int u, v, w, p;
};

Edge E[MAXN+10];

int K;
int cnt;

ll dp1[MAXN+10], dp2[MAXN+10];

void dfs(int now, int bef)
{
	sz[now]=1;
	dfn[now]=++cnt;
	for(auto nxt : adj[now])
	{
		if(nxt.first==bef) continue;
		par[nxt.first]=now;
		val[nxt.first]=nxt.second;
		sz[now]+=sz[nxt.first];
		dfs(nxt.first, now);
		dp1[now]+=dp2[nxt.first];
	}
	dp2[now]=dp1[now]+E[val[now]].w;
}

struct Node
{
	ll sum, cnt;
	Node *lc, *rc;
	Node() : sum(0), cnt(0), lc(NULL), rc(NULL) {}
};

struct SEG
{
	Node *root;

	SEG()
	{
		root=new Node();
	}

	void push(Node *node, ll tl, ll tr, ll k)
	{
		node->sum+=k;
		node->cnt++;
		if(tl==tr) return;
		ll mid=(tl+tr-1)/2;
		if(k<=mid)
		{
			if(node->lc==NULL) node->lc=new Node();
			push(node->lc, tl, mid, k);
		}
		else
		{
			if(node->rc==NULL) node->rc=new Node();
			push(node->rc, mid+1, tr, k);
		}	
	}

	void pop(Node *node, ll tl, ll tr, ll k)
	{
		node->sum-=k;
		node->cnt--;
		if(tl==tr) return;
		ll mid=(tl+tr-1)/2;
		if(k<=mid)
		{
			if(node->lc==NULL) assert(0);
			pop(node->lc, tl, mid, k);
		}
		else
		{
			if(node->rc==NULL) assert(0);
			pop(node->rc, mid+1, tr, k);
		}
	}

	ll query(Node *node, ll tl, ll tr, ll k)
	{
		if(k<=0) return 0;
		if(node->cnt==0) return 0;
		if(tl==tr)
		{
			return tl*min(k, node->cnt);
		}
		ll mid=(tl+tr-1)/2;
		if(node->lc==NULL) return query(node->rc, mid+1, tr, k);
		if(node->rc==NULL) return query(node->lc, tl, mid, k);

		if(node->lc->cnt>=k)
		{
			return query(node->lc, tl, mid, k);
		}
		else
		{
			return query(node->rc, mid+1, tr, k-node->lc->cnt)+node->lc->sum;
		}
	}

	void push(ll k)
	{
		if(k>=0) return;
		push(root, -1e15, 0, k);
	}

	void pop(ll k)
	{
		if(k>=0) return;
		pop(root, -1e15, 0, k);
	}

	ll query(ll k)
	{
		return query(root, -1e15, 0, k);
	}
};

SEG seg[MAXN+10];
ll sum[MAXN+10];

int chk[MAXN+10];
int chk2[MAXN+10];

bool vis[MAXN+10];

void dfs2(int now, int bef)
{
	vis[now]=true;
	while(!adj2[now].empty() && min(deg[E[adj2[now].back().second].u], deg[E[adj2[now].back().second].v])<=K)
	{
		//printf("??POP %d %d\n", now, adj2[now].back().first);
		adj2[now].pop_back();
	}

	if(chk[par[now]]==0)
	{
		sum[par[now]]-=dp2[now];
		seg[par[now]].pop(dp1[now]-dp2[now]);
	}

	for(auto nxt : adj2[now])
	{
		if(vis[nxt.first]) continue;
		if(nxt.first==bef) continue;
		dfs2(nxt.first, now);
	}

	dp1[now]=sum[now]+seg[now].query(K-1);
	dp2[now]=sum[now]+seg[now].query(K)+E[val[now]].w;
	//printf("%d : %lld %lld\n", now, dp1[now], dp2[now]);

	if(chk[par[now]]==0)
	{
		sum[par[now]]+=dp2[now];
		seg[par[now]].push(dp1[now]-dp2[now]);
	}
}

vector<ll> minimum_closure_costs(int _N, vector<int> _U, vector<int> _V, vector<int> _W)
{
	N=_N;
	for(int i=1; i<N; i++)
	{
		int u=_U[i-1]+1, v=_V[i-1]+1, w=_W[i-1];
		adj[u].push_back({v, i});
		adj[v].push_back({u, i});
		adj2[u].push_back({v, i});
		adj2[v].push_back({u, i});
		E[i]={u, v, w, i};
		deg[u]++;
		deg[v]++;
	}
	for(int i=1; i<=N; i++)
	{
		sort(adj2[i].begin(), adj2[i].end(), [&](const pii &p, const pii &q)
		{
			return min(deg[E[p.second].u], deg[E[p.second].v])>min(deg[E[q.second].u], deg[E[q.second].v]);
		});
	}


	dfs(1, 1);
	ans[0]=dp1[1];
	for(int i=1; i<=N; i++)
	{
		sum[par[i]]+=dp2[i];
		//printf("HIHI\n");
		seg[par[i]].push(dp1[i]-dp2[i]);
	}

	vector<pii> V;
	for(int i=1; i<=N; i++) V.push_back({deg[i], i});
	sort(V.begin(), V.end());
	for(int i=1, j=0; i<N; i++)
	{
		K=i;
		//printf("K %d\n", K);
		for(; j<V.size() && V[j].first<=K; j++)
		{
			int u=V[j].second;
			chk[u]++;
			if(chk[par[u]]==0)
			{
				sum[par[u]]-=dp2[u];
				seg[par[u]].pop(dp1[u]-dp2[u]);
				sum[par[u]]+=E[val[u]].w;
				seg[par[u]].push(-E[val[u]].w);
			}
			for(auto it : adj[u])
			{
				chk2[it.second]++;
			}
		}

		vector<int> VV;
		for(int k=j; k<V.size(); k++)
		{
			VV.push_back(V[k].second);
		}
		sort(VV.begin(), VV.end(), [&](const int &p, const int &q)
		{
			return dfn[p]<dfn[q];
		});
		for(int u : VV)
		{
			if(vis[u]) continue;
			//printf("HI %d\n", u);
			dfs2(u, u);
			ans[K]+=min(dp1[u], dp2[u]);
		}
		for(int k=j; k<V.size(); k++)
		{
			int u=V[k].second;
			vis[u]=false;
		}
	}

	return vector<ll>(ans, ans+N);
}

Compilation message (stderr)

roads.cpp: In function 'std::vector<long long int> minimum_closure_costs(int, std::vector<int>, std::vector<int>, std::vector<int>)':
roads.cpp:216:10: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  216 |   for(; j<V.size() && V[j].first<=K; j++)
      |         ~^~~~~~~~~
roads.cpp:234:17: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  234 |   for(int k=j; k<V.size(); k++)
      |                ~^~~~~~~~~
roads.cpp:249:17: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  249 |   for(int k=j; k<V.size(); k++)
      |                ~^~~~~~~~~

• 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