Submission #1113889

#	Time	Username	Problem	Language	Result	Execution time	Memory
1113889		Thunnus	Toll (APIO13_toll)	C++17	100 / 100	1117 ms	15048 KiB

toll

//Challenge: Accepted
#include <iostream>
#include <algorithm>
#include <vector>
#include <utility>
#include <bitset>
#include <set>
#include <queue>
#include <stack>
#include <assert.h>
#include <cmath>
#include <iomanip>
#include <random>
using namespace std;
void debug(){cout << endl;};
template<class T, class ...U> void debug(T a, U ... b){cout << a << " ", debug(b ...);};
template<class T> void pary(T l, T r) {
	while (l != r) cout << *l << " ", l++;
	cout << endl;
};
#define ll long long
#define maxn 100005
#define maxc 25
#define mod 1000000007
#define pii pair<int, int>
#define ff first
#define ss second
#define io ios_base::sync_with_stdio(0);cin.tie(0);
/*
Basic idea:
Consider the brute force solution of enumerating all 2^k subsets of new edges chosen. 
We add these edges first and maintain DSU for connectivity, then add original edges by increasing weight.
When there is a cycle, we can update the new roads' maximum weight for it to be chosen.
Finally, we do a DFS to calculate the total revenue. Complexity O(mlogm + m*2^k)
 
To speed this up, notice that exactly x edges from the original set are chosen to "limit" the x new edges' weights.
If we try building a spanning tree for all edges (original and new), there will be <= k edges that were in the original
tree but might get replaced. It can be proven that we only need to consider these edges.
Thus, we construct a "compressed graph" as follows: we partition the original graph's vertices by their component in the 
combined MST if we remove the new edges. Then we can use the basic idea above on the compressed graph and the limiting edges
to update the answer. Notice that there might be cycles in the original k edges, however the conclusion remains the same.
Time Complexity: O(mlogm + 2^k * k ^2)
*/
struct edge {
	int u, v, w;
	edge(){u = v = w = 0;}
	edge(int x, int y, int z){u = x, v = y, w = z;}
};
vector<edge> ed;
struct DSU{
	int par[maxn];
	void init(int n) {
		for (int i = 0;i <= n;i++) par[i] = i;
	}
	int find(int a) {
		return a == par[a] ? a : (par[a] = find(par[a]));
	}
	void Union(int a, int b) { // set par[a] = b
		par[find(a)] = find(b);
	}
} dsu, comp, built;
 
bool mark[maxn];
int vis[maxn];
int wei[maxc], pa[maxc], dep[maxc]; // for the compressed tree, wei[i]:maximum cost possible for edge (i, pa[i])
int c[maxn]; 
ll p[maxc], siz[maxc]; //siz[i]: number of people passing through edge (i, pa[i]) in compressed tree
vector<pii> small[maxc]; //compressed tree adjacency list
void dfs(int n, int par, int d, ll &ans) { //builds compressed tree sizes
	siz[n] = p[n];
	pa[n] = par;
	dep[n] = d;
	for (auto v:small[n]) {
		if (v.ff != par) {
			dfs(v.ff, n, d+1, ans);
			siz[n] += siz[v.ff];
			if (v.ss) ans += siz[v.ff] * wei[v.ff];
		}
	}
}
void addedge(int a, int b, int val) { //updates added edges
	if (dep[a] < dep[b]) swap(a, b);
	while (dep[a] > dep[b]) {
		wei[a] = min(wei[a], val);
		a = pa[a];
	}	
	while (a != b) {
		wei[a] = min(wei[a], val), wei[b] = min(wei[b], val);
		a = pa[a], b = pa[b];
	}
}
 
int main() {
	io
	int n, m, k;
	cin >> n >> m >> k;
	for (int i = 0;i < m;i++) {
		int u, v, w;
		cin >> u >> v >> w;
		ed.push_back(edge(u, v, w));
	}
	sort(ed.begin(), ed.end(), [&](edge x, edge y){return x.w < y.w;});
	dsu.init(n);
	built.init(n);
	vector<edge> add, tree, rep; //add: new k edges, tree: edges in MST, rep: edges that might be used for updating cost
	for (int i = 0;i < k;i++) {
		int u, v;
		cin >> u >> v;
		mark[u] = mark[v] = 1;
		built.Union(u, v);
		add.push_back(edge(u, v, 0));
	}
	for (int i = 1;i <= n;i++) cin >> c[i];
	for (auto e:ed) {
		if (dsu.find(e.u) != dsu.find(e.v)) {
			dsu.Union(e.u, e.v);
			tree.push_back(e);
		}	
	}
	dsu.init(n);
	int ind = 0;
	for (auto e:tree) { //groups vertices and finds edges in rep
		e.u = dsu.find(e.u), e.v = dsu.find(e.v);
		if (mark[e.u] && mark[e.v]) {
			if (built.find(e.u) == built.find(e.v)) rep.push_back(e); 
			else {
				dsu.Union(e.u, e.v);
			}
		} else if (mark[e.u] || mark[e.v]) {
			if (mark[e.u]) swap(e.u, e.v);
			dsu.Union(e.u, e.v);
		} else {
			dsu.Union(e.u, e.v);
		}
		built.Union(e.u, e.v);
	}
	for (int i = 1;i <= n;i++) { //relabels vertices [0 ... ind-1]
		int f = dsu.find(i);
		if (!vis[f]) {
			vis[f] = ++ind;
		}
		p[vis[f]-1] += c[i];
	}	
	//relabels edges
	for (auto &e:add) e.u = vis[dsu.find(e.u)]-1, e.v = vis[dsu.find(e.v)]-1;
	for (auto &e:rep) e.u = vis[dsu.find(e.u)]-1, e.v = vis[dsu.find(e.v)]-1;
 
	ll ans = 0;
	assert(rep.size() <= k);
	for (int i = 1;i < (1<<k);i++) {
		for (int j = 0;j < ind;j++) small[j].clear(), wei[j] = 1<<30;
		comp.init(ind); //comp: compressed graph DSU
		bool cy = 0;
		for (int j = 0;j < k;j++) {
			if (i & (1<<j)) {
				small[add[j].u].push_back({add[j].v, 1});
				small[add[j].v].push_back({add[j].u, 1});	
				if (comp.find(add[j].u) == comp.find(add[j].v)) {
					cy = 1;
					break;
				}
				comp.Union(add[j].u, add[j].v);
			}
		}
		if (cy) continue; //stop if new edges form a cycle
		vector<edge> lim;
		for (auto e:rep) {
			if (comp.find(e.u) == comp.find(e.v)) {
				lim.push_back(e); //found limiting edge
			} else {
				comp.Union(e.u, e.v);
				small[e.u].push_back({e.v, 0});
				small[e.v].push_back({e.u, 0});
			}
		}	
		ll tmp = 0;
		dfs(0, 0, 0, tmp);
		for (auto e:lim) addedge(e.u, e.v, e.w);
		tmp = 0;
		dfs(0, 0, 0, tmp);
		ans = max(ans, tmp);
	}	
	cout << ans << endl;
}

Compilation message (stderr)

In file included from toll.cpp:10:
toll.cpp: In function 'int main()':
toll.cpp:149:20: warning: comparison of integer expressions of different signedness: 'std::vector<edge>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
  149 |  assert(rep.size() <= k);
      |         ~~~~~~~~~~~^~~~

• Subtask 1: 16 / 16
• Subtask 2: 18 / 18
• Subtask 3: 22 / 22
• Subtask 4: 22 / 22
• Subtask 5: 22 / 22