shichikuji and power grid

shichikuji.cpp

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

typedef long long ll;
const int MAXN = 2e3 + 10;

ll n, x[MAXN], y[MAXN], c[MAXN], k[MAXN];
int pai[MAXN], sze[MAXN]; 

vector<tuple<ll, int, int>> arestas; // peso, vertice 1, vertice 2
		   
int find(int u)
{
	if(pai[u] == u) return u;
	return pai[u] = find(pai[u]);
}

bool join(int u, int v) // retorna true se juntei, false caso contrario
{ 
	u = find(u), v = find(v);
	if(u == v) return false; // ja estao na mesma componente

	if(sze[u] < sze[v]) swap(u, v);
	pai[v] = u; sze[u] += sze[v];
	return true;
}

ll calcCusto(int i, int j)
{
	ll dist = abs(x[i] - x[j]) + abs(y[i] - y[j]);
	return (k[i] + k[j]) * dist;
}

int main()
{
 	ios_base::sync_with_stdio(false);
	cin.tie(NULL);

	cin >> n;

	for(int i = 1; i <= n; i++)
		cin >> x[i] >> y[i];

	for(int i = 1; i <= n; i++) cin >> c[i];
	for(int i = 1; i <= n; i++) cin >> k[i];

	//Arestas entre cidades
	for(int i = 1; i <= n; i++)
		for(int j = i + 1; j <= n; j++)
			arestas.emplace_back(calcCusto(i, j), i, j);

	//Aresta com vertice ficticio
	for(int i = 1; i <= n; i++) 
		arestas.emplace_back(c[i], i, 0);

	sort(arestas.begin(), arestas.end());

	for(int i = 0; i <= n; i++) // Inicializo a DSU
		sze[i] = 1, pai[i] = i;

	ll custoTotal = 0;
	vector<int> estacoes;
	vector<pair<int, int>> ligacoes;	

	for(auto [w, u, v] : arestas){
		if(!join(u, v)) continue; // Nao esta na MST
		custoTotal += w;

		if(v == 0) estacoes.push_back(u); 
		else ligacoes.emplace_back(u, v);
	}

	cout << custoTotal << '\n';
	cout << (int) estacoes.size() << '\n';

	for(auto cidade : estacoes)
		cout << cidade << ' ';

	cout << '\n' << (int) ligacoes.size() << '\n';

	for(auto [u, v] : ligacoes)
		cout << u << ' ' << v << '\n';
}